Skip to content

Commit

Permalink
implement iterator for Bound iterator
Browse files Browse the repository at this point in the history
  • Loading branch information
davidhewitt committed Dec 24, 2023
1 parent 7d24584 commit e5ab4fc
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 12 deletions.
6 changes: 6 additions & 0 deletions src/ffi_ptr_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand Down
25 changes: 19 additions & 6 deletions src/instance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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, or NULL.
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,
Expand Down
71 changes: 65 additions & 6 deletions src/types/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,28 @@ impl PyIterator {
}
}

/// Implementation of functionality for [`PyIterator`].
///
/// These methods are defined for the `Bound<'py, PyIterator>` smart pointer, so to use method call
/// syntax these methods are separated into a trait, because stable Rust does not yet support
/// `arbitrary_self_types`.
#[doc(alias = "PyIterator")]
pub trait PyIteratorMethods<'py> {
/// Equivalent of `next(self)`.
fn next(&self) -> Option<PyResult<Bound<'py, PyAny>>>;
}

impl<'py> PyIteratorMethods<'py> for Bound<'py, PyIterator> {
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<'p> Iterator for &'p PyIterator {
type Item = PyResult<&'p PyAny>;

Expand All @@ -56,21 +78,58 @@ 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();
Bound::borrowed_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> {
PyIteratorMethods::next(self)
}

#[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> 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> {
PyIteratorMethods::next(*self)
}

#[cfg(not(Py_LIMITED_API))]
fn size_hint(&self) -> (usize, Option<usize>) {
(*self).size_hint()
}
}

impl PyTypeCheck for PyIterator {
const NAME: &'static str = "Iterator";

Expand Down

0 comments on commit e5ab4fc

Please sign in to comment.