Merge pull request #268 from kngwyu/fix-dict-iter

Fix PyDictIterator's segfault(for #159)

examples/rustapi_module/setup.py

@@ -60,6 +60,11 @@ setup(
             "Cargo.toml",
             rustc_flags=get_py_version_cfgs(),
         ),
+        RustExtension(
+            "rustapi_module.test_dict",
+            "Cargo.toml",
+            rustc_flags=get_py_version_cfgs(),
+        ),
     ],
     install_requires=install_requires,
     tests_require=tests_require,

examples/rustapi_module/src/dict_iter.rs

@@ -0,0 +1,43 @@
+use pyo3::prelude::*;
+
+use pyo3::exceptions::RuntimeError;
+use pyo3::types::PyDict;
+
+#[pymodinit(test_dict)]
+fn test_dict(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<DictSize>()?;
+    Ok(())
+}
+
+#[pyclass]
+pub struct DictSize {
+    expected: u32,
+}
+
+#[pymethods]
+impl DictSize {
+    #[new]
+    fn __new__(obj: &PyRawObject, expected: u32) -> PyResult<()> {
+        obj.init(|| DictSize { expected })
+    }
+
+    fn iter_dict(&mut self, _py: Python, dict: &PyDict) -> PyResult<u32> {
+        let mut seen = 0u32;
+        for (sym, values) in dict.iter() {
+            seen += 1;
+            println!(
+                "{:4}/{:4} iterations:{}=>{}",
+                seen, self.expected, sym, values
+            );
+        }
+
+        if seen == self.expected {
+            Ok(seen)
+        } else {
+            Err(PyErr::new::<RuntimeError, _>(format!(
+                "Expected {} iterations - performed {}",
+                self.expected, seen
+            )))
+        }
+    }
+}

examples/rustapi_module/src/lib.rs

@@ -4,5 +4,6 @@
 extern crate pyo3;
 
 pub mod datetime;
+pub mod dict_iter;
 pub mod othermod;
 pub mod subclassing;

examples/rustapi_module/tests/test_dict_iter.py

@@ -0,0 +1,13 @@
+from rustapi_module.test_dict import DictSize
+import pytest
+
+
+@pytest.mark.parametrize(
+    "size",
+    [64, 128, 256],
+)
+def test_size(size):
+    d = {}
+    for i in range(size):
+        d[i] = str(i)
+    assert DictSize(len(d)).iter_dict(d) == size

src/types/dict.rs

@@ -148,17 +148,23 @@ impl PyDict {
     /// Note that it's unsafe to use when the dictionary might be changed
     /// by other python code.
     pub fn iter(&self) -> PyDictIterator {
-        PyDictIterator { dict: self, pos: 0 }
+        let py = self.py();
+        PyDictIterator {
+            dict: self.to_object(py),
+            pos: 0,
+            py,
+        }
     }
 }
 
-pub struct PyDictIterator<'a> {
+pub struct PyDictIterator<'py> {
-    dict: &'a PyDict,
+    dict: PyObject,
     pos: isize,
+    py: Python<'py>,
 }
 
-impl<'a> Iterator for PyDictIterator<'a> {
+impl<'py> Iterator for PyDictIterator<'py> {
-    type Item = (&'a PyObjectRef, &'a PyObjectRef);
+    type Item = (&'py PyObjectRef, &'py PyObjectRef);
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
@@ -166,7 +172,7 @@ impl<'a> Iterator for PyDictIterator<'a> {
             let mut key: *mut ffi::PyObject = mem::uninitialized();
             let mut value: *mut ffi::PyObject = mem::uninitialized();
             if ffi::PyDict_Next(self.dict.as_ptr(), &mut self.pos, &mut key, &mut value) != 0 {
-                let py = self.dict.py();
+                let py = self.py;
                 Some((py.from_borrowed_ptr(key), py.from_borrowed_ptr(value)))
             } else {
                 None