Preserve element order when initializing Python dictionary (#57)

* Change dictionary literal initializer implementation * Add test for dictionary initializer * Fix CI tests
pvieito · Sep 1, 2022 · 060e1c8 · 060e1c8
1 parent 5105e31
commit 060e1c8
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 11 deletions.
diff --git a/PythonKit/Python.swift b/PythonKit/Python.swift
@@ -1274,8 +1274,8 @@ extension PythonObject : SignedNumeric {
         return self < 0 ? -self : self
     }
 
-    //override the default implementation of - prefix function
-    //from SignedNumeric  (https://bugs.swift.org/browse/SR-13293)
+    // Override the default implementation of `-` prefix function
+    // from SignedNumeric (https://bugs.swift.org/browse/SR-13293).
     public static prefix func - (_ operand: Self) -> Self {
         return performUnaryOp(PyNumber_Negative, operand: operand)
     }
@@ -1472,8 +1472,36 @@ extension PythonObject : ExpressibleByArrayLiteral, ExpressibleByDictionaryLiter
     }
     public typealias Key = PythonObject
     public typealias Value = PythonObject
+
+    // Preserves element order in the final Python object, unlike
+    // `Dictionary.pythonObject`. When keys are duplicated, throw the same
+    // runtime error as `Swift.Dictionary.init(dictionaryLiteral:)`. This
+    // differs from Python's key uniquing semantics, which silently override an
+    // existing key with the next one it encounters.
     public init(dictionaryLiteral elements: (PythonObject, PythonObject)...) {
-        self.init(Dictionary(elements, uniquingKeysWith: { lhs, _ in lhs }))
+        _ = Python // Ensure Python is initialized.
+        let dict = PyDict_New()!
+        for (key, value) in elements {
+            let k = key.ownedPyObject
+            let v = value.ownedPyObject
+
+            // Use Python's native key checking instead of querying whether
+            // `elements` contains the key. Although this could theoretically
+            // produce different results, it produces the Python object we want.
+            switch PyDict_Contains(dict, k) {
+            case 0:
+                PyDict_SetItem(dict, k, v)
+            case 1:
+                fatalError("Dictionary literal contains duplicate keys")
+            default:
+                try! throwPythonErrorIfPresent()
+                fatalError("No result or error checking whether \(elements) contains \(key)")
+            }
+
+            Py_DecRef(k)
+            Py_DecRef(v)
+        }
+        self.init(consuming: dict)
     }
 }
 
@@ -1876,7 +1904,9 @@ public struct PythonClass {
                 (key, value.pythonObject)
             }
 
-            dictionary = Dictionary(castedElements, uniquingKeysWith: { lhs, _ in lhs })
+            dictionary = Dictionary(castedElements, uniquingKeysWith: { _, _ in
+                fatalError("Dictionary literal contains duplicate keys")
+            })
         }
     }
 

diff --git a/PythonKit/PythonLibrary+Symbols.swift b/PythonKit/PythonLibrary+Symbols.swift
@@ -93,6 +93,10 @@ let PyErr_Fetch: @convention(c) (
 let PyDict_New: @convention(c) () -> PyObjectPointer? =
     PythonLibrary.loadSymbol(name: "PyDict_New")
 
+let PyDict_Contains: @convention(c) (
+    PyObjectPointer?, PyObjectPointer?) -> Int32 =
+    PythonLibrary.loadSymbol(name: "PyDict_Contains")
+
 let PyDict_SetItem: @convention(c) (
     PyObjectPointer?, PyObjectPointer, PyObjectPointer) -> Void =
     PythonLibrary.loadSymbol(name: "PyDict_SetItem")

diff --git a/Tests/PythonKitTests/PythonRuntimeTests.swift b/Tests/PythonKitTests/PythonRuntimeTests.swift
@@ -34,7 +34,7 @@ class PythonRuntimeTests: XCTestCase {
         XCTAssertEqual(2, Python.len(dict))
         XCTAssertEqual(1, dict["a"])
         XCTAssertEqual(0.5, dict[1])
-
+        
         XCTAssertEqual(2, dict.count as Int)
         XCTAssertEqual(2, dict.checking.count!)
         XCTAssertEqual(2, dict.throwing.count!)
@@ -43,6 +43,33 @@ class PythonRuntimeTests: XCTestCase {
         XCTAssertEqual("c", dict["b"])
         dict["b"] = "d"
         XCTAssertEqual("d", dict["b"])
+
+        // Dictionary initializer patch does not work on Python 2, but that
+        // version is no longer being actively supported.
+        guard Python.versionInfo.major >= 3 else {
+            return
+        }
+
+        // Pandas DataFrame regression test spotted in Jupyter. This is
+        // non-deterministic, so repeat it several times to ensure the bug does
+        // not appear.
+        for _ in 0..<15 {
+            let records: [PythonObject] = [
+                ["col 1": 3, "col 2": 5],
+                ["col 1": 8, "col 2": 2]
+            ]
+            let records_description =
+                "[{'col 1': 3, 'col 2': 5}, {'col 1': 8, 'col 2': 2}]"
+            XCTAssertEqual(String(describing: records), records_description)
+
+            let records_alt: [PythonObject] = [
+                ["col 1": 3, "col 2": 5, "col 3": 4],
+                ["col 1": 8, "col 2": 2, "col 3": 4]
+            ]
+            let records_alt_description =
+                "[{'col 1': 3, 'col 2': 5, 'col 3': 4}, {'col 1': 8, 'col 2': 2, 'col 3': 4}]"
+            XCTAssertEqual(String(describing: records_alt), records_alt_description)
+        }
     }
 
     func testRange() {
@@ -123,12 +150,12 @@ class PythonRuntimeTests: XCTestCase {
     }
 
     func testUnaryOps() {
-	var x = PythonObject(5)
-	x = -x
-	XCTAssertEqual(-5, x)
-	x = PythonObject(-5)
-	x = -x
-	XCTAssertEqual(5, x)
+        var x = PythonObject(5)
+        x = -x
+        XCTAssertEqual(-5, x)
+        x = PythonObject(-5)
+        x = -x
+        XCTAssertEqual(5, x)
     }
 
     func testComparable() {