Subscript Dictionary with String-based Enums in Swift

Question

I want to extend Dictionary with String keys (JSON dictionaries) to allow subscripting with any enum that has a RawValue type of String. The end goal would be multiple enums that can be used to subscript JSON dictionaries.

enum JSONKey: String {
    case one, two, three
}

enum OtherJSONKey: String {
    case a, b, c
}

if let one = jsonDictionary[.one] { /* ... */ }
if let b = jsonDictionary[.b] { /* ... */ }

But I can't figure out how to implement this. I know I need to extend Dictionary, but can't figure out either the generic extension constraints or method extension constraints.

My first idea was to try to add a generic constraint to the subscript method. I don't think that subscript methods allow generics, though.

extension Dictionary {
    subscript<T: RawRepresentable>(key: T) -> Value? { /* ... */ }
}

Even if putting generic constraints on a subscript worked, I still need a way to nest my generic constraints. Or to constrain the dictionary to keys that are string-based enums. To put it in code that isn't valid, I want to do this:

extension Dictionary where Key: RawRepresentable where RawValue == String {
    subscript(key: Key) -> Value { /* ... */ }
}

// or

extension Dictionary {
    subscript<T: RawRepresentable where RawValue == String>(key: T) -> Value { /* ... */ }
}

Is extending Dictionary to accept string-based enums as a subscript actually possible?

My other thoughts on how to implement something like this included enum inheritance and creating a protocol for particular enums that I want to use as subscripts. I know some of this can't be done, but figured it was worth mentioning the idea. So, again, to put it in code that isn't valid:

enum JSONKey: String {}
enum NumbersJSONKey: JSONKey {
    case one, two, three
}
enum LettersJSONKey: JSONKey {
    case a, b, c
}

// or

protocol JSONKeys {}
enum NumbersJSONKey: JSONKey {
    case one, two, three
}
enum LettersJSONKey: JSONKey {
    case a, b, c
}

// then subscript with
if let one = json[.one] { /* ... */ }

Update:

I have played with this some more and gotten a little closer. The extension below compiles, but gives me a "subscript is ambiguous" error if I actually try to use it.

extension Collection where Iterator.Element == (key: String, value: AnyObject) {

    // Compiles but can't be used because of ambiguous subscript.
    subscript(key: CustomStringConvertible) -> AnyObject? {
        guard let i = index(where: { $0.key == key.description }) else { return nil }
        return self[i].value
    }

}

@titaniumdecoy's answer works so it will be the accepted answer unless someone can come up with something better.

Answer 1

With Swift 4's support for Generic Subscripts you can now do this:

extension Dictionary where Key: ExpressibleByStringLiteral {
    subscript<Index: RawRepresentable>(index: Index) -> Value? where Index.RawValue == String {
        get {
            return self[index.rawValue as! Key]
        }

        set {
            self[index.rawValue as! Key] = newValue
        }
    }
} 

Which allows you to use any Enum which has String as it's RawValue type:

let value = jsonDict[JSONKey.one]

This will now work for any String Enum, not just JSONKey

Answer 2

As I understand it, you want an extension on any Dictionary with String keys to allow subscripting using an enum with String as its RawValue type. If so, the following should work for you:

enum JSONKey: String {
    case one, two, three
}

class JSONObject { }

extension Dictionary where Key: StringLiteralConvertible {
    subscript(jsonKey: JSONKey) -> Value? {
        get {
            return self[jsonKey.rawValue as! Key]
        }
        set {
            self[jsonKey.rawValue as! Key] = newValue
        }
    }
}

var jsonDict: [String: AnyObject] = [:]    

jsonDict[JSONKey.one] = JSONObject()
jsonDict["two"] = JSONObject()

print(jsonDict["one"]!)
print(jsonDict[JSONKey.two]!)

---

If you want to extend this to work for any enum with String as its RawValue type, you need generics. As Swift does not support generic subscripts (see SR-115), get/set methods or a property would be required:

enum AnotherEnum: String {
    case anotherCase
}

extension Dictionary where Key: StringLiteralConvertible {
    func getValue<T: RawRepresentable where T.RawValue == String>(forKey key: T) -> Value? {
        return self[key.rawValue as! Key]
    }
    mutating func setValue<T: RawRepresentable where T.RawValue == String>(value: Value, forKey key: T) {
        self[key.rawValue as! Key] = value
    }
}

jsonDict.setValue(JSONObject(), forKey: AnotherEnum.anotherCase)
print(jsonDict.getValue(forKey: AnotherEnum.anotherCase)!)