Remove nested key from dictionary

Question

Let\'s say I have a rather complex dictionary, like this one:

let dict: [String: Any] = [
    \"countries\": [
        \"japan\": [
            \"capital\":          


        

Answer 1

I'd to like to follow up on my previous answer with another solution. This one extends Swift's Dictionary type with a new subscript that takes a key path.

I first introduce a new type named KeyPath to represent a key path. It's not strictly necessary, but it makes working with key paths much easier because it lets us wrap the logic of splitting a key path into its components.

import Foundation

/// Represents a key path.
/// Can be initialized with a string of the form "this.is.a.keypath"
///
/// We can't use Swift's #keyPath syntax because it checks at compilet time
/// if the key path exists.
struct KeyPath {
    var elements: [String]

    var isEmpty: Bool { return elements.isEmpty }
    var count: Int { return elements.count }
    var path: String {
        return elements.joined(separator: ".")
    }

    func headAndTail() -> (String, KeyPath)? {
        guard !isEmpty else { return nil }
        var tail = elements
        let head = tail.removeFirst()
        return (head, KeyPath(elements: tail))
    }
}

extension KeyPath {
    init(_ string: String) {
        elements = string.components(separatedBy: ".")
    }
}

extension KeyPath: ExpressibleByStringLiteral {
    init(stringLiteral value: String) {
        self.init(value)
    }
    init(unicodeScalarLiteral value: String) {
        self.init(value)
    }
    init(extendedGraphemeClusterLiteral value: String) {
        self.init(value)
    }
}

Next I create a dummy protocol named StringProtocol that we later need to constrain our Dictionary extension. Swift 3.0 doesn't yet support extensions on generic types that constrain a generic parameter to a concrete type (such as extension Dictionary where Key == String). Support for this is planned for Swift 4.0, but until then, we need this little workaround:

// We need this because Swift 3.0 doesn't support extension Dictionary where Key == String
protocol StringProtocol {
    init(string s: String)
}

extension String: StringProtocol {
    init(string s: String) {
        self = s
    }
}

Now we can write the new subscripts. The implementation for the getter and setter are fairly long, but they should be straightforward: we traverse the key path from beginning to end and then get/set the value at that position:

// We want extension Dictionary where Key == String, but that's not supported yet,
// so work around it with Key: StringProtocol.
extension Dictionary where Key: StringProtocol {
    subscript(keyPath keyPath: KeyPath) -> Any? {
        get {
            guard let (head, remainingKeyPath) = keyPath.headAndTail() else {
                return nil
            }

            let key = Key(string: head)
            let value = self[key]
            switch remainingKeyPath.isEmpty {
            case true:
                // Reached the end of the key path
                return value
            case false:
                // Key path has a tail we need to traverse
                switch value {
                case let nestedDict as [Key: Any]:
                    // Next nest level is a dictionary
                    return nestedDict[keyPath: remainingKeyPath]
                default:
                    // Next nest level isn't a dictionary: invalid key path, abort
                    return nil
                }
            }
        }
        set {
            guard let (head, remainingKeyPath) = keyPath.headAndTail() else {
                return
            }
            let key = Key(string: head)

            // Assign new value if we reached the end of the key path
            guard !remainingKeyPath.isEmpty else {
                self[key] = newValue as? Value
                return
            }

            let value = self[key]
            switch value {
            case var nestedDict as [Key: Any]:
                // Key path has a tail we need to traverse
                nestedDict[keyPath: remainingKeyPath] = newValue
                self[key] = nestedDict as? Value
            default:
                // Invalid keyPath
                return
            }
        }
    }
}

And this is how it looks in use:

var dict: [String: Any] = [
    "countries": [
        "japan": [
            "capital": [
                "name": "tokyo",
                "lat": "35.6895",
                "lon": "139.6917"
            ],
            "language": "japanese"
        ]
    ],
    "airports": [
        "germany": ["FRA", "MUC", "HAM", "TXL"]
    ]
]

dict[keyPath: "countries.japan"] // ["language": "japanese", "capital": ["lat": "35.6895", "name": "tokyo", "lon": "139.6917"]]
dict[keyPath: "countries.someothercountry"] // nil
dict[keyPath: "countries.japan.capital"] // ["lat": "35.6895", "name": "tokyo", "lon": "139.6917"]
dict[keyPath: "countries.japan.capital.name"] // "tokyo"
dict[keyPath: "countries.japan.capital.name"] = "Edo"
dict[keyPath: "countries.japan.capital.name"] // "Edo"
dict[keyPath: "countries.japan.capital"] // ["lat": "35.6895", "name": "Edo", "lon": "139.6917"]

I really like this solution. It's quite a lot of code, but you only have to write it once and I think it looks very nice in use.