Skip to content

Commit

Permalink
implement PyIterator without additional lifetime
Browse files Browse the repository at this point in the history
This lets us treat it no different from other types
like PySequence.
  • Loading branch information
birkenfeld committed Sep 11, 2020
1 parent a3d378b commit 9364080
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 28 deletions.
2 changes: 1 addition & 1 deletion src/types/any.rs
Original file line number Diff line number Diff line change
Expand Up @@ -362,7 +362,7 @@ impl PyAny {
///
/// This is typically a new iterator but if the argument is an iterator,
/// this returns itself.
pub fn iter(&self) -> PyResult<PyIterator> {
pub fn iter(&self) -> PyResult<&PyIterator> {
PyIterator::from_object(self.py(), self)
}

Expand Down
58 changes: 31 additions & 27 deletions src/types/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,12 @@
//
// based on Daniel Grunwald's https://github.com/dgrunwald/rust-cpython

use crate::{ffi, AsPyPointer, PyAny, PyErr, PyNativeType, PyResult, Python};
use crate::{
ffi, AsPyPointer, PyAny, PyDowncastError, PyErr, PyNativeType, PyResult, PyTryFrom, Python,
};

/// A Python iterator object.
///
/// Unlike other Python objects, this class includes a `Python<'p>` token
/// so that `PyIterator` can implement the Rust `Iterator` trait.
///
/// This means that you can't use `PyIterator` in many places where other
/// types like `PyList` can automatically be extracted from objects, such
/// as function arguments.
///
/// # Example
///
/// ```rust
Expand All @@ -30,31 +25,24 @@ use crate::{ffi, AsPyPointer, PyAny, PyErr, PyNativeType, PyResult, Python};
/// # }
/// ```
#[derive(Debug)]
pub struct PyIterator<'p>(&'p PyAny);
#[repr(transparent)]
pub struct PyIterator(PyAny);
pyobject_native_type_named!(PyIterator);
pyobject_native_type_extract!(PyIterator);

impl<'p> PyIterator<'p> {
impl PyIterator {
/// Constructs a `PyIterator` from a Python iterable object.
///
/// Equivalent to Python's built-in `iter` function.
pub fn from_object<T>(py: Python<'p>, obj: &T) -> PyResult<PyIterator<'p>>
pub fn from_object<'p, T>(py: Python<'p>, obj: &T) -> PyResult<&'p PyIterator>
where
T: AsPyPointer,
{
let iter = unsafe {
// This looks suspicious, but is actually correct. Even though ptr is an owned
// reference, PyIterator takes ownership of the reference and decreases the count
// in its Drop implementation.
//
// Therefore we must use from_borrowed_ptr_or_err instead of from_owned_ptr_or_err so
// that the GILPool does not take ownership of the reference.
py.from_borrowed_ptr_or_err(ffi::PyObject_GetIter(obj.as_ptr()))?
};

Ok(PyIterator(iter))
unsafe { py.from_owned_ptr_or_err(ffi::PyObject_GetIter(obj.as_ptr())) }
}
}

impl<'p> Iterator for PyIterator<'p> {
impl<'p> Iterator for &'p PyIterator {
type Item = PyResult<&'p PyAny>;

/// Retrieves the next item from an iterator.
Expand All @@ -79,10 +67,26 @@ impl<'p> Iterator for PyIterator<'p> {
}
}

/// Dropping a `PyIterator` instance decrements the reference count on the object by 1.
impl<'p> Drop for PyIterator<'p> {
fn drop(&mut self) {
unsafe { ffi::Py_DECREF(self.0.as_ptr()) }
impl<'v> PyTryFrom<'v> for PyIterator {
fn try_from<V: Into<&'v PyAny>>(value: V) -> Result<&'v PyIterator, PyDowncastError<'v>> {
let value = value.into();
unsafe {
if ffi::PyIter_Check(value.as_ptr()) != 0 {
Ok(<PyIterator as PyTryFrom>::try_from_unchecked(value))
} else {
Err(PyDowncastError::new(value, "Iterator"))
}
}
}

fn try_from_exact<V: Into<&'v PyAny>>(value: V) -> Result<&'v PyIterator, PyDowncastError<'v>> {
<PyIterator as PyTryFrom>::try_from(value)
}

#[inline]
unsafe fn try_from_unchecked<V: Into<&'v PyAny>>(value: V) -> &'v PyIterator {
let ptr = value.into() as *const _ as *const PyIterator;
&*ptr
}
}

Expand Down

0 comments on commit 9364080

Please sign in to comment.