From d669a943f73b67574890595d57c3f1fffc71ad51 Mon Sep 17 00:00:00 2001
From: David Hewitt <mail@davidhewitt.dev>
Date: Sat, 23 Dec 2023 14:26:02 +0000
Subject: [PATCH] implement iterator for `Bound` iterator

---
 src/ffi_ptr_ext.rs    |  6 ++++++
 src/instance.rs       | 25 +++++++++++++++++++------
 src/types/iterator.rs | 43 +++++++++++++++++++++++++++++++++++++------
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/src/ffi_ptr_ext.rs b/src/ffi_ptr_ext.rs
index 0eca2c36dd5..8a45f5d70c0 100644
--- a/src/ffi_ptr_ext.rs
+++ b/src/ffi_ptr_ext.rs
@@ -16,6 +16,7 @@ use sealed::Sealed;
 
 pub(crate) trait FfiPtrExt: Sealed {
     unsafe fn assume_owned_or_err(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>>;
+    unsafe fn assume_owned_or_opt(self, py: Python<'_>) -> Option<Bound<'_, PyAny>>;
     unsafe fn assume_owned(self, py: Python<'_>) -> Bound<'_, PyAny>;
 
     /// Assumes this pointer is borrowed from a parent object.
@@ -41,6 +42,11 @@ impl FfiPtrExt for *mut ffi::PyObject {
         Bound::from_owned_ptr_or_err(py, self)
     }
 
+    #[inline]
+    unsafe fn assume_owned_or_opt(self, py: Python<'_>) -> Option<Bound<'_, PyAny>> {
+        Bound::from_owned_ptr_or_opt(py, self)
+    }
+
     #[inline]
     unsafe fn assume_owned(self, py: Python<'_>) -> Bound<'_, PyAny> {
         Bound::from_owned_ptr(py, self)
diff --git a/src/instance.rs b/src/instance.rs
index 8406e86abbe..0bd55b895e1 100644
--- a/src/instance.rs
+++ b/src/instance.rs
@@ -49,18 +49,31 @@ pub struct Bound<'py, T>(Python<'py>, ManuallyDrop<Py<T>>);
 
 impl<'py> Bound<'py, PyAny> {
     /// Constructs a new Bound from a pointer. Panics if ptr is null.
+    ///
+    /// # Safety
+    ///
+    /// `ptr`` must be a valid pointer to a Python object.
     pub(crate) unsafe fn from_owned_ptr(py: Python<'py>, ptr: *mut ffi::PyObject) -> Self {
         Self(py, ManuallyDrop::new(Py::from_owned_ptr(py, ptr)))
     }
 
-    // /// Constructs a new Bound from a pointer. Returns None if ptr is null.
-    // ///
-    // /// Safety: ptr must be a valid pointer to a Python object, or NULL.
-    // pub unsafe fn from_owned_ptr_or_opt(py: Python<'py>, ptr: *mut ffi::PyObject) -> Option<Self> {
-    //     Py::from_owned_ptr_or_opt(py, ptr).map(|obj| Self(py, ManuallyDrop::new(obj)))
-    // }
+    /// Constructs a new Bound from a pointer. Returns None if ptr is null.
+    ///
+    /// # Safety
+    ///
+    /// `ptr`` must be a valid pointer to a Python object, or NULL.
+    pub(crate) unsafe fn from_owned_ptr_or_opt(
+        py: Python<'py>,
+        ptr: *mut ffi::PyObject,
+    ) -> Option<Self> {
+        Py::from_owned_ptr_or_opt(py, ptr).map(|obj| Self(py, ManuallyDrop::new(obj)))
+    }
 
     /// Constructs a new Bound from a pointer. Returns error if ptr is null.
+    ///
+    /// # Safety
+    ///
+    /// `ptr` must be a valid pointer to a Python object, or NULL.
     pub(crate) unsafe fn from_owned_ptr_or_err(
         py: Python<'py>,
         ptr: *mut ffi::PyObject,
diff --git a/src/types/iterator.rs b/src/types/iterator.rs
index bf392398eb9..0e033232964 100644
--- a/src/types/iterator.rs
+++ b/src/types/iterator.rs
@@ -1,4 +1,5 @@
 use crate::ffi_ptr_ext::FfiPtrExt;
+use crate::instance::Borrowed;
 use crate::py_result_ext::PyResultExt;
 use crate::{
     ffi, AsPyPointer, Bound, PyAny, PyDowncastError, PyErr, PyNativeType, PyResult, PyTypeCheck,
@@ -56,21 +57,51 @@ impl<'p> Iterator for &'p PyIterator {
     /// Further `next()` calls after an exception occurs are likely
     /// to repeatedly result in the same exception.
     fn next(&mut self) -> Option<Self::Item> {
-        let py = self.0.py();
+        Borrowed::<PyIterator>::from_gil_ref(self)
+            .next()
+            .map(|result| result.map(Bound::into_gil_ref))
+    }
 
-        match unsafe { py.from_owned_ptr_or_opt(ffi::PyIter_Next(self.0.as_ptr())) } {
-            Some(obj) => Some(Ok(obj)),
-            None => PyErr::take(py).map(Err),
-        }
+    #[cfg(not(Py_LIMITED_API))]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        Bound::borrowed_from_gil_ref(self).size_hint()
+    }
+}
+
+impl<'py> Iterator for Bound<'py, PyIterator> {
+    type Item = PyResult<Bound<'py, PyAny>>;
+
+    /// Retrieves the next item from an iterator.
+    ///
+    /// Returns `None` when the iterator is exhausted.
+    /// If an exception occurs, returns `Some(Err(..))`.
+    /// Further `next()` calls after an exception occurs are likely
+    /// to repeatedly result in the same exception.
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        Borrowed::from(&*self).next()
     }
 
     #[cfg(not(Py_LIMITED_API))]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let hint = unsafe { ffi::PyObject_LengthHint(self.0.as_ptr(), 0) };
+        let hint = unsafe { ffi::PyObject_LengthHint(self.as_ptr(), 0) };
         (hint.max(0) as usize, None)
     }
 }
 
+impl<'py> Borrowed<'_, 'py, PyIterator> {
+    // TODO: this method is on Borrowed so that &'py PyIterator can use this; once that
+    // implementation is deleted this method should be moved to the `Bound<'py, PyIterator> impl
+    fn next(self) -> Option<PyResult<Bound<'py, PyAny>>> {
+        let py = self.py();
+
+        match unsafe { ffi::PyIter_Next(self.as_ptr()).assume_owned_or_opt(py) } {
+            Some(obj) => Some(Ok(obj)),
+            None => PyErr::take(py).map(Err),
+        }
+    }
+}
+
 impl PyTypeCheck for PyIterator {
     const NAME: &'static str = "Iterator";