Add parsing of dictionary/collection mirrors

adam-fowler · Jun 7, 2021 · 9095158 · 9095158
1 parent ddc3ae2
commit 9095158
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 12 deletions.
diff --git a/Sources/JMESPath/Variable.swift b/Sources/JMESPath/Variable.swift
@@ -34,27 +34,41 @@ public enum JMESVariable {
             self = .array(set.map { $0 })
         case let dictionary as [String: Any]:
             self = .object(dictionary)
+        case is NSNull:
+            self = .null
         default:
-            if any is NSNull {
-                self = .null
-                return
-            }
             // use Mirror to build JMESVariable.object
             let mirror = Mirror(reflecting: any)
             guard mirror.children.count > 0 else {
                 self = .other(any)
                 return
             }
-            var object: JMESObject = [:]
-            for child in mirror.children {
-                guard let label = child.label else {
-                    self = .null
-                    return
+            switch mirror.displayStyle {
+            case .collection:
+                let array = mirror.children.map {
+                    Self.unwrap($0.value) ?? NSNull()
+                }
+                self = .array(array)
+            case .dictionary:
+                var object: JMESObject = [:]
+                var index: Int = 0
+                while let key = mirror.descendant(index, "key") as? String, let value = mirror.descendant(index, "value") {
+                    object[key] = Self.unwrap(value) ?? NSNull()
+                    index += 1
+                }
+                self = .object(object)
+            default:
+                var object: JMESObject = [:]
+                for child in mirror.children {
+                    guard let label = child.label else {
+                        self = .null
+                        return
+                    }
+                    let unwrapValue = Self.unwrap(child.value) ?? NSNull()
+                    object[label] = unwrapValue
                 }
-                let unwrapValue = Self.unwrap(child.value) ?? NSNull()
-                object[label] = unwrapValue
+                self = .object(object)
             }
-            self = .object(object)
         }
     }
 

diff --git a/Tests/JMESPathTests/MirrorTests.swift b/Tests/JMESPathTests/MirrorTests.swift
@@ -75,4 +75,22 @@ final class MirrorTests: XCTestCase {
         let test = TestObject(e: .test2)
         self.testInterpreter("e", data: test, result: TestEnum.test2)
     }
+
+    func testCustomReflectableArray() {
+        struct TestObject: CustomReflectable {
+            let a: [Int]
+            var customMirror: Mirror { return Mirror(reflecting: a) }
+        }
+        let test = TestObject(a: [1,2,3,4])
+        self.testInterpreter("[2]", data: test, result: 3)
+    }
+
+    func testCustomReflectableDictionary() {
+        struct TestObject: CustomReflectable {
+            let d: [String: String]
+            var customMirror: Mirror { return Mirror(reflecting: d) }
+        }
+        let test = TestObject(d: ["test": "one", "test2": "two", "test3": "three"])
+        self.testInterpreter("test2", data: test, result: "two")
+    }
 }