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": [
                "name": "tokyo",
                "lat": "35.6895",
                "lon": "139.6917"
            ],
            "language": "japanese"
        ]
    ],
    "airports": [
        "germany": ["FRA", "MUC", "HAM", "TXL"]
    ]
]

I can access all fields with if let .. blocks optionally casting to something that I can work with, when reading.

However, I am currently writing unit tests where I need to selectively break dictionaries in multiple ways.

But I don't know how to elegantly remove keys from the dictionary.

For example I want to remove the key "japan" in one test, in the next "lat" should be nil.

Here's my current implementation for removing "lat":

if var countries = dict["countries"] as? [String: Any],
    var japan = countries["japan"] as? [String: Any],
    var capital = japan["capital"] as? [String: Any]
    {
        capital.removeValue(forKey: "lat")
        japan["capital"] = capital
        countries["japan"] = japan
        dictWithoutLat["countries"] = countries
}

Surely there must be a more elegant way?

Ideally I'd write a test helper that takes a KVC string and has a signature like:

func dictWithoutKeyPath(_ path: String) -> [String: Any] 

In the "lat" case I'd call it with dictWithoutKeyPath("countries.japan.capital.lat").

Answer 1

When working with a subscript, if the subscript is get/set and the variable is mutable, then the entire expression is mutable. However, due to the type cast the expression "loses" the mutability. (It's not an l-value anymore).

The shortest way to solve this is by creating a subscript that is get/set and does the conversion for you.

extension Dictionary {
    subscript(jsonDict key: Key) -> [String:Any]? {
        get {
            return self[key] as? [String:Any]
        }
        set {
            self[key] = newValue as? Value
        }
    }
}

Now you can write the following:

dict[jsonDict: "countries"]?[jsonDict: "japan"]?[jsonDict: "capital"]?["name"] = "berlin"

We liked this question so much that we decided to make a (public) Swift Talk episode about it: mutating untyped dictionaries

Answer 2

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.