From 34099b33f85e8fd826f09b2386e14fbdcc833ef4 Mon Sep 17 00:00:00 2001
From: kngwyu <yuji.kngw.80s.revive@gmail.com>
Date: Tue, 13 Nov 2018 00:24:59 +0900
Subject: [PATCH] Fix PyDictIterator's segfault(for #159)

---
 examples/rustapi_module/setup.py              |  5 +++
 examples/rustapi_module/src/dict_iter.rs      | 42 +++++++++++++++++++
 examples/rustapi_module/src/lib.rs            |  1 +
 .../rustapi_module/tests/test_dict_iter.py    | 12 ++++++
 src/types/dict.rs                             | 16 ++++---
 5 files changed, 71 insertions(+), 5 deletions(-)
 create mode 100644 examples/rustapi_module/src/dict_iter.rs
 create mode 100644 examples/rustapi_module/tests/test_dict_iter.py

diff --git a/examples/rustapi_module/setup.py b/examples/rustapi_module/setup.py
index 7be4bdea..1b9ed7f6 100644
--- a/examples/rustapi_module/setup.py
+++ b/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,
diff --git a/examples/rustapi_module/src/dict_iter.rs b/examples/rustapi_module/src/dict_iter.rs
new file mode 100644
index 00000000..7394b156
--- /dev/null
+++ b/examples/rustapi_module/src/dict_iter.rs
@@ -0,0 +1,42 @@
+use pyo3::prelude::*;
+
+use pyo3::exceptions as exc;
+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(|_t| 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
+            );
+        }
+
+        match seen == self.expected {
+            true => Ok(seen),
+            _ => Err(PyErr::new::<exc::TypeError, _>(format!(
+                "Expected {} iterations - performed {}",
+                self.expected, seen
+            ))),
+        }
+    }
+}
diff --git a/examples/rustapi_module/src/lib.rs b/examples/rustapi_module/src/lib.rs
index f0b66eb6..36f8a7b3 100644
--- a/examples/rustapi_module/src/lib.rs
+++ b/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;
diff --git a/examples/rustapi_module/tests/test_dict_iter.py b/examples/rustapi_module/tests/test_dict_iter.py
new file mode 100644
index 00000000..b5878867
--- /dev/null
+++ b/examples/rustapi_module/tests/test_dict_iter.py
@@ -0,0 +1,12 @@
+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(0,size):
+        d[i] = str(i)
+    assert DictSize(len(d)).iter_dict(d) == size
diff --git a/src/types/dict.rs b/src/types/dict.rs
index 00da0bdf..22d4da23 100644
--- a/src/types/dict.rs
+++ b/src/types/dict.rs
@@ -8,7 +8,7 @@ use crate::object::PyObject;
 use crate::python::{IntoPyPointer, Python, ToPyPointer};
 use crate::types::{PyList, PyObjectRef};
 use std;
-use std::{cmp, collections, hash, mem};
+use std::{cmp, collections, hash, marker::PhantomData, mem};
 
 /// Represents a Python `dict`.
 #[repr(transparent)]
@@ -148,13 +148,19 @@ 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,
+            __marker: PhantomData,
+        }
     }
 }
 
-pub struct PyDictIterator<'a> {
-    dict: &'a PyDict,
+pub struct PyDictIterator<'dict> {
+    dict: PyObject,
     pos: isize,
+    __marker: PhantomData<&'dict ()>,
 }
 
 impl<'a> Iterator for PyDictIterator<'a> {
@@ -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 = Python::assume_gil_acquired();
                 Some((py.from_borrowed_ptr(key), py.from_borrowed_ptr(value)))
             } else {
                 None