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

Interesting question. The problem seems to be that Swift's optional chaining mechanism, which is normally capable of mutating nested dictionaries, trips over the necessary type casts from Any to [String:Any]. So while accessing a nested element becomes only unreadable (because of the typecasts):

// E.g. Accessing countries.japan.capital
((dict["countries"] as? [String:Any])?["japan"] as? [String:Any])?["capital"]

… mutating a nested element doesn't even work:

// Want to mutate countries.japan.capital.name.
// The typecasts destroy the mutating optional chaining.
((((dict["countries"] as? [String:Any])?["japan"] as? [String:Any])?["capital"] as? [String:Any])?["name"] as? String) = "Edo"
// Error: Cannot assign to immutable expression

---

Possible Solution

The idea is to get rid of the untyped dictionary and convert it into a strongly typed structure where each element has the same type. I admit that this is a heavy-handed solution, but it works quite well in the end.

An enum with associated values would work well for our custom type that replaces the untyped dictionary:

enum KeyValueStore {
    case dict([String: KeyValueStore])
    case array([KeyValueStore])
    case string(String)
    // Add more cases for Double, Int, etc.
}

The enum has one case for each expected element type. The three cases cover your example, but it could easily be expanded to cover more types.

Next, we define two subscripts, one for keyed access to a dictionary (with strings) and one for indexed access to an array (with integers). The subscripts check if self is a .dict or .array respectively and if so return the value at the given key/index. They return nil if the type doesn't match, e.g. if you tried to access a key of a .string value. The subscripts also have setters. This is key to make chained mutation work:

extension KeyValueStore {
    subscript(_ key: String) -> KeyValueStore? {
        // If self is a .dict, return the value at key, otherwise return nil.
        get {
            switch self {
            case .dict(let d):
                return d[key]
            default:
                return nil
            }
        }
        // If self is a .dict, mutate the value at key, otherwise ignore.
        set {
            switch self {
            case .dict(var d):
                d[key] = newValue
                self = .dict(d)
            default:
                break
            }
        }
    }

    subscript(_ index: Int) -> KeyValueStore? {
        // If self is an array, return the element at index, otherwise return nil.
        get {
            switch self {
            case .array(let a):
                return a[index]
            default:
                return nil
            }
        }
        // If self is an array, mutate the element at index, otherwise return nil.
        set {
            switch self {
            case .array(var a):
                if let v = newValue {
                    a[index] = v
                } else {
                    a.remove(at: index)
                }
                self = .array(a)
            default:
                break
            }
        }
    }
}

Lastly, we add some convenience initializers for initializing our type with dictionary, array or string literals. These are not strictly necessary, but make working with the type easier:

extension KeyValueStore: ExpressibleByDictionaryLiteral {
    init(dictionaryLiteral elements: (String, KeyValueStore)...) {
        var dict: [String: KeyValueStore] = [:]
        for (key, value) in elements {
            dict[key] = value
        }
        self = .dict(dict)
    }
}

extension KeyValueStore: ExpressibleByArrayLiteral {
    init(arrayLiteral elements: KeyValueStore...) {
        self = .array(elements)
    }
}

extension KeyValueStore: ExpressibleByStringLiteral {
    init(stringLiteral value: String) {
        self = .string(value)
    }

    init(extendedGraphemeClusterLiteral value: String) {
        self = .string(value)
    }

    init(unicodeScalarLiteral value: String) {
        self = .string(value)
    }
}

And here's the example:

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

// Now optional chaining works:
keyValueStore["countries"]?["japan"]?["capital"]?["name"] // .some(.string("tokyo"))
keyValueStore["countries"]?["japan"]?["capital"]?["name"] = "Edo"
keyValueStore["countries"]?["japan"]?["capital"]?["name"] // .some(.string("Edo"))
keyValueStore["airports"]?["germany"]?[1] // .some(.string("MUC"))
keyValueStore["airports"]?["germany"]?[1] = "BER"
keyValueStore["airports"]?["germany"]?[1] // .some(.string("BER"))
// Remove value from array by assigning nil. I'm not sure if this makes sense.
keyValueStore["airports"]?["germany"]?[1] = nil
keyValueStore["airports"]?["germany"] // .some(array([.string("FRA"), .string("HAM"), .string("TXL")]))

Answer 2

Pass your dictionary to this function, It will return you a flat dictionary, without any nested dict incorporated .

//SWIFT 3.0

func concatDict( dict: [String: Any])-> [String: Any]{
        var dict = dict
        for (parentKey, parentValue) in dict{
            if let insideDict = parentValue as? [String: Any]{
                let keys = insideDict.keys.map{
                    return parentKey + $0
                }
                for (key, value) in zip(keys, insideDict.values) {
                    dict[key] = value
                }
                dict[parentKey] = nil
                dict = concatDict(dict: dict)
            }
        }
        return dict
    }

Answer 3

You could construct recursive methods (read/write) which visits your given key path by repeatedly attempting conversions of (sub-)dictionary values to [Key: Any] dictionaries themselves. Moreover, allow public access to these methods via a new subscript.

Note that you might have to explicitly import Foundation for access to the components(separatedBy:) method of String (bridged).

extension Dictionary {       
    subscript(keyPath keyPath: String) -> Any? {
        get {
            guard let keyPath = Dictionary.keyPathKeys(forKeyPath: keyPath) 
                else { return nil }
            return getValue(forKeyPath: keyPath)
        }
        set {
            guard let keyPath = Dictionary.keyPathKeys(forKeyPath: keyPath),
                let newValue = newValue else { return }
            self.setValue(newValue, forKeyPath: keyPath)
        }
    }

    static private func keyPathKeys(forKeyPath: String) -> [Key]? {
        let keys = forKeyPath.components(separatedBy: ".")
            .reversed().flatMap({ $0 as? Key })
        return keys.isEmpty ? nil : keys
    }

    // recursively (attempt to) access queried subdictionaries
    // (keyPath will never be empty here; the explicit unwrapping is safe)
    private func getValue(forKeyPath keyPath: [Key]) -> Any? {
        guard let value = self[keyPath.last!] else { return nil }
        return keyPath.count == 1 ? value : (value as? [Key: Any])
                .flatMap { $0.getValue(forKeyPath: Array(keyPath.dropLast())) }
    }

    // recursively (attempt to) access the queried subdictionaries to
    // finally replace the "inner value", given that the key path is valid
    private mutating func setValue(_ value: Any, forKeyPath keyPath: [Key]) {
        guard self[keyPath.last!] != nil else { return }            
        if keyPath.count == 1 {
            (value as? Value).map { self[keyPath.last!] = $0 }
        }
        else if var subDict = self[keyPath.last!] as? [Key: Value] {
            subDict.setValue(value, forKeyPath: Array(keyPath.dropLast()))
            (subDict as? Value).map { self[keyPath.last!] = $0 }
        }
    }
}

Example setup

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

Example usage:

// read value for a given key path
let isNil: Any = "nil"
print(dict[keyPath: "countries.japan.capital.name"] ?? isNil) // tokyo
print(dict[keyPath: "airports"] ?? isNil)                     // ["germany": ["FRA", "MUC", "HAM", "TXL"]]
print(dict[keyPath: "this.is.not.a.valid.key.path"] ?? isNil) // nil

// write value for a given key path
dict[keyPath: "countries.japan.language"] = "nihongo"
print(dict[keyPath: "countries.japan.language"] ?? isNil) // nihongo

dict[keyPath: "airports.germany"] = 
    (dict[keyPath: "airports.germany"] as? [Any] ?? []) + ["FOO"]
dict[keyPath: "this.is.not.a.valid.key.path"] = "notAdded"

print(dict)
/*  [
        "countries": [
            "japan": [
                "capital": [
                    "name": "tokyo", 
                    "lon": "139.6917",
                    "lat": "35.6895"
                    ], 
                "language": "nihongo"
            ]
        ], 
        "airports": [
            "germany": ["FRA", "MUC", "HAM", "TXL", "FOO"]
        ]
    ] */

Note that if a supplied key path does not exist for an assignment (using setter), this will not result in the construction of the equivalent nested dictionary, but simply result in no mutation of the dictionary.

Answer 4

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 5

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.