Default to _PyDict_Next() iterator

ijl · Jul 31, 2022 · 8f47042 · 8f47042
1 parent 3374311
commit 8f47042
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 26 deletions.
diff --git a/build.rs b/build.rs
@@ -7,12 +7,11 @@ fn main() {
     println!("cargo:rerun-if-env-changed=CFLAGS");
     println!("cargo:rerun-if-env-changed=LDFLAGS");
     println!("cargo:rerun-if-env-changed=RUSTFLAGS");
-    println!("cargo:rerun-if-env-changed=ORJSON_DISABLE_PYDICTITER");
+    println!("cargo:rerun-if-env-changed=ORJSON_ENABLE_PYDICTITER");
     println!("cargo:rerun-if-env-changed=ORJSON_DISABLE_YYJSON");
 
     let py_cfg = pyo3_build_config::get();
     py_cfg.emit_pyo3_cfgs();
-    let py_version_minor = py_cfg.version.minor;
 
     if let Some(true) = version_check::supports_feature("core_intrinsics") {
         println!("cargo:rustc-cfg=feature=\"intrinsics\"");
@@ -22,7 +21,7 @@ fn main() {
         println!("cargo:rustc-cfg=feature=\"optimize\"");
     }
 
-    if std::env::var("ORJSON_DISABLE_PYDICTITER").is_err() && py_version_minor < 11 {
+    if std::env::var("ORJSON_ENABLE_PYDICTITER").is_ok() {
         println!("cargo:rustc-cfg=feature=\"pydictiter\"");
     }
 

diff --git a/src/ffi/dict.rs b/src/ffi/dict.rs
@@ -126,7 +126,7 @@ impl Iterator for PyDictIter {
                 while self.idx < self.len {
                     self.idx += 1;
                     if !(*entry_ptr).me_value.is_null() {
-                        return Some(((*entry_ptr).me_key, (*entry_ptr).me_value))
+                        return Some(((*entry_ptr).me_key, (*entry_ptr).me_value));
                     }
                     entry_ptr = entry_ptr.add(1);
                 }

diff --git a/src/serialize/dataclass.rs b/src/serialize/dataclass.rs
@@ -9,7 +9,6 @@ use crate::unicode::*;
 use crate::ffi::PyDictIter;
 use serde::ser::{Serialize, SerializeMap, Serializer};
 
-use std::ptr::addr_of_mut;
 use std::ptr::NonNull;
 
 pub struct DataclassFastSerializer {
@@ -49,26 +48,7 @@ impl Serialize for DataclassFastSerializer {
             return serializer.serialize_map(Some(0)).unwrap().end();
         }
         let mut map = serializer.serialize_map(None).unwrap();
-        let mut pos = 0isize;
-        let mut key: *mut pyo3_ffi::PyObject = std::ptr::null_mut();
-        let mut value: *mut pyo3_ffi::PyObject = std::ptr::null_mut();
-        for _ in 0..=len.saturating_sub(1) {
-            unsafe {
-                pyo3_ffi::_PyDict_Next(
-                    self.ptr,
-                    addr_of_mut!(pos),
-                    addr_of_mut!(key),
-                    addr_of_mut!(value),
-                    std::ptr::null_mut(),
-                )
-            };
-            let pyvalue = PyObjectSerializer::new(
-                value,
-                self.opts,
-                self.default_calls,
-                self.recursion + 1,
-                self.default,
-            );
+        for (key, value) in PyDictIter::from_pyobject(self.ptr) {
             if unlikely!(unsafe { ob_type!(key) != STR_TYPE }) {
                 err!(SerializeError::KeyMustBeStr)
             }
@@ -80,6 +60,13 @@ impl Serialize for DataclassFastSerializer {
             if unlikely!(key_as_str.as_bytes()[0] == b'_') {
                 continue;
             }
+            let pyvalue = PyObjectSerializer::new(
+                value,
+                self.opts,
+                self.default_calls,
+                self.recursion + 1,
+                self.default,
+            );
             map.serialize_key(key_as_str).unwrap();
             map.serialize_value(&pyvalue)?;
         }

diff --git a/test/test_dict.py b/test/test_dict.py
@@ -24,7 +24,7 @@ def test_dict_pop(self):
         """Test pop and replace a key in a dict with no other keys."""
         data = {"id": "any"}
         data.pop("id")
-        assert orjson.dumps(data) == b'{}'
+        assert orjson.dumps(data) == b"{}"
         data["id"] = "new"
         assert orjson.dumps(data) == b'{"id":"new"}'
 
@@ -49,3 +49,11 @@ def test_dict_0xffff(self):
         data["112"] = 1
         data["113"] = 2
         assert orjson.loads(orjson.dumps(data)) == data
+
+    def test_dict_dict(self):
+        class C:
+            def __init__(self):
+                self.a = 0
+                self.b = 1
+
+        assert orjson.dumps(C().__dict__) == b'{"a":0,"b":1}'