Adds key mapping support for use in destination plugins. (#38)

* Adds key mapping support for use in destination plugins.

* Slight cleanup.

* Added array processing to tests.

Co-authored-by: Brandon Sneed <[email protected]>

Author: bsneed

Sources/Segment/Utilities/JSON.swift

@@ -267,9 +267,73 @@ extension JSON {
     }
 }
 
+// MARK: - Mapping
+
+extension JSON {
+    /// Maps keys supplied, in the format of ["Old": "New"].  Gives an optional value transformer that can be used to transform values based on the final key name.
+    /// - Parameters:
+    ///   - keys: A dictionary containing key mappings, in the format of ["Old": "New"].
+    ///   - valueTransform: An optional value transform closure.  Key represents the new key name.
+    public func mapKeys(_ keys: [String: String], valueTransform: ((_ key: String, _ value: Any) -> Any)? = nil) throws -> JSON {
+        guard let dict = self.dictionaryValue else { return self }
+        let mapped = try dict.mapKeys(keys, valueTransform: valueTransform)
+        let result = try JSON(mapped)
+        return result
+    }
+}
 
 // MARK: - Helpers
 
+extension Dictionary where Key == String, Value == Any {
+    internal func mapKeys(_ keys: [String: String], valueTransform: ((_ key: Key, _ value: Value) -> Any)? = nil) throws -> [Key: Value] {
+        let mapped = Dictionary(uniqueKeysWithValues: self.map { key, value -> (Key, Value) in
+            var newKey = key
+            var newValue = value
+            
+            // does this key have a mapping?
+            if keys.keys.contains(key) {
+                if let mappedKey = keys[key] {
+                    // if so, lets change the key to the new value.
+                    newKey = mappedKey
+                }
+            }
+            // is this value a dictionary?
+            if let dictValue = value as? [Key: Value] {
+                if let r = try? dictValue.mapKeys(keys, valueTransform: valueTransform) {
+                    // if so, lets recurse...
+                    newValue = r
+                }
+            } else if let arrayValue = value as? [Value] {
+                // if it's an array, we need to see if any dictionaries are within and process
+                // those as well.
+                newValue = arrayValue.map { item -> Value in
+                    var newValue = item
+                    if let dictValue = item as? [Key: Value] {
+                        if let r = try? dictValue.mapKeys(keys, valueTransform: valueTransform) {
+                            newValue = r
+                        }
+                    }
+                    return newValue
+                }
+            }
+            
+            if !(newValue is [Key: Value]), let transform = valueTransform {
+                // it's not a dictionary apply our transform.
+                
+                // note: if it's an array, we've processed any dictionaries inside
+                // already, but this gives the opportunity to apply a transform to the other
+                // items in the array that weren't dictionaries.
+                newValue = transform(newKey, newValue)
+            }
+            
+            return (newKey, newValue)
+        })
+        
+        return mapped
+    }
+}
+
+
 fileprivate extension NSNumber {
     static let trueValue = NSNumber(value: true)
     static let trueObjCType = trueValue.objCType

Tests/Segment-Tests/JSON_Tests.swift

@@ -130,4 +130,51 @@ class JSONTests: XCTestCase {
 
         XCTAssertTrue(fetchedTraits?.email == "[email protected]")
     }
+    
+    func testKeyMapping() {
+        let keys = ["Key1": "AKey1", "Key2": "AKey2"]
+        let dict: [String: Any] = ["Key1": 1, "Key2": 2, "Key3": 3, "Key4": ["Key1": 1]]
+        
+        let json = try! JSON(dict)
+        
+        let output = try! json.mapKeys(keys).dictionaryValue
+        
+        XCTAssertTrue(output!["AKey1"] as! Int == 1)
+        XCTAssertTrue(output!["AKey2"] as! Int == 2)
+        XCTAssertTrue(output!["Key3"] as! Int == 3)
+        
+        let subDict = output!["Key4"] as! [String: Any]
+        XCTAssertTrue(subDict["AKey1"] as! Int == 1)
+    }
+    
+    func testKeyMappingWithValueTransform() {
+        let keys = ["Key1": "AKey1", "Key2": "AKey2"]
+        let dict: [String: Any] = ["Key1": 1, "Key2": 2, "Key3": 3, "Key4": ["Key1": 1], "Key5": [1, 2, ["Key1": 1]]]
+        
+        let json = try! JSON(dict)
+        
+        let output = try! json.mapKeys(keys, valueTransform: { key, value in
+            var newValue = value
+            if let v = newValue as? Int {
+                if v == 1 {
+                    newValue = 11
+                }
+            }
+            print("value = \(value.self)")
+            return newValue
+        }).dictionaryValue
+        
+        XCTAssertTrue(output!["AKey1"] as! Int == 11)
+        XCTAssertTrue(output!["AKey2"] as! Int == 2)
+        XCTAssertTrue(output!["Key3"] as! Int == 3)
+        
+        let subDict = output!["Key4"] as! [String: Any]
+        XCTAssertTrue(subDict["AKey1"] as! Int == 11)
+        
+        let subArray = output!["Key5"] as! [Any]
+        let subArrayDict = subArray[2] as! [String: Any]
+        XCTAssertTrue(subArray[0] as! Int == 1)
+        XCTAssertTrue(subArray[1] as! Int == 2)
+        XCTAssertTrue(subArrayDict["AKey1"] as! Int == 11)
+    }
 }