Swift: Reading and writing to (possible) nested dictionaries for a given key path, using a recursive approach

DictionaryKeyPath.swift

// Inspired by: https://gist.github.com/dfrib/d7419038f7e680d3f268750d63f0dfae
import Foundation

extension Dictionary {
    subscript(keyPath string: String) -> Value? {
        get {
            return self[keyPath: Dictionary.keyPath(for: string)]
        }
        set {
            self[keyPath: Dictionary.keyPath(for: string)] = newValue
        }
    }
    
    subscript(keyPath keyPath: Key...) -> Value? {
        get {
            return self[keyPath: keyPath]
        }
        set {
            self[keyPath: keyPath] = newValue
        }
    }
    
    subscript(keyPath keyPath: [Key]) -> Value? {
        get {
            guard !keyPath.isEmpty else { return nil }
            return getValue(forKeyPath: keyPath)
        }
        set {
            guard !keyPath.isEmpty else { return }
            self.setValue(newValue, forKeyPath: keyPath)
            if newValue == nil {
                cleanUp(forKeyPath: keyPath)
            }
        }
    }
    
    static private func keyPath(for keyPathString: String) -> [Key] {
        let keys = keyPathString.components(separatedBy: "/")
        return keys.compactMap({ $0 as? Key })
    }
    
    // recursively (attempt to) access queried subdictionaries
    // (keyPath will never be empty here; the explicit unwrapping is safe)
    private func getValue(forKeyPath keyPath: [Key]) -> Value? {
        guard let value = self[keyPath.first!] else { return nil }
        return keyPath.count == 1 ? value : (value as? [Key: Value])
            .flatMap { $0.getValue(forKeyPath: Array(keyPath.dropFirst())) }
    }
    
    // recursively access, create or overwrite the
    // queried subdictionaries to finally set the "inner value"
    private mutating func setValue(_ value: Value?, forKeyPath keyPath: [Key]) {
        if keyPath.count == 1 {
            self[keyPath.first!] = value
        } else {
            var subDict = self[keyPath.first!] as? [Key: Value] ?? [Key: Value]()
            subDict.setValue(value, forKeyPath: Array(keyPath.dropFirst()))
            self[keyPath.first!] = subDict as? Value
        }
    }
    
    // recursively (attempt to) remove left over empty subdictionaries
    private mutating func cleanUp(forKeyPath keyPath: [Key]) {
        // Never set root to nil
        guard !keyPath.isEmpty else { return }
        
        if let value = getValue(forKeyPath: keyPath) {
            guard let dict = value as? [Key: Value], dict.isEmpty else {
                // This endpoint does not continue cleanUp because
                // a non nil value that isn't an empty dict is found
                return
            }
            setValue(nil, forKeyPath: keyPath)
        }
        cleanUp(forKeyPath: Array(keyPath.dropLast()))
    }
}

/* ------------------------------------------------------------------ */
// example usage
var dict: [String: Any] = [
  "some": [
    "nested": [
      "data": "Often found in json"
    ]
  ]
]

let a = dict[keyPath: "some/nested/data"]         // "Often found in json"
let b = dict[keyPath: "some", "nested", "data"]   // "Often found in json"
let c = dict[keyPath: ["some", "nested", "data"]] // "Often found in json"

dict[keyPath: "some/nested/data"] = "Replacing data"
dict[keyPath: "another/nested/path"] = "Creating new subdirectories"
print(dict as AnyObject)
/*
 {
     another =     {
         nested =         {
            path = "Creating new subdirectories";
         };
     };
     some =     {
         nested =         {
            data = "Replacing data";
         };
     };
 }
 */

// Clean up removes all empty subdirectories
dict[keyPath: "some/nested/data"] = nil
print(dict as AnyObject)
/*
 {
     another =     {
         nested =         {
            path = "Creating new subdirectories";
         };
     };
 }
 */

// Warning: I opted to allow overwriting of data
dict[keyPath: "another/nested/path/new/path"] = "Overwrote path node"
print(dict as AnyObject)
/*
 {
     another =     {
         nested =         {
             path =             {
                 new =                 {
                    path = "Overwrote path node";
                 };
             };
         };
     };
 }
 */

// This also means if we set a deeper node to nil
// cleanup will delete higher empty levels since
// those are overwritten
dict[keyPath: "another/nested/path/new/path/even/deeper/path"] = nil
print(dict as AnyObject)
/*
 {
 }
 */

Future updates are now at https://gist.github.com/Amzd/e908d4bf3cbad2c9d766586baa619566