Merge pull request #997 from davidhewitt/pyproto-next-return

Add __next__ return

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - Add `PyByteArray::data`, `PyByteArray::as_bytes`, and `PyByteArray::as_bytes_mut`. [#967](https://github.com/PyO3/pyo3/pull/967)
 - Add `GILOnceCell` to use in situations where `lazy_static` or `once_cell` can deadlock. [#975](https://github.com/PyO3/pyo3/pull/975)
 - Add `Py::borrow`, `Py::borrow_mut`, `Py::try_borrow`, and `Py::try_borrow_mut` for accessing `#[pyclass]` values. [#976](https://github.com/PyO3/pyo3/pull/976)
+- Add `IterNextOutput` and `IterANextOutput` for returning from `__next__` / `__anext__`. [#997](https://github.com/PyO3/pyo3/pull/997)
 
 ### Changed
 - Simplify internals of `#[pyo3(get)]` attribute. (Remove the hidden API `GetPropertyValue`.) [#934](https://github.com/PyO3/pyo3/pull/934)

examples/rustapi_module/setup.py

@@ -99,6 +99,7 @@ setup(
         make_rust_extension("rustapi_module.othermod"),
         make_rust_extension("rustapi_module.subclassing"),
         make_rust_extension("rustapi_module.test_dict"),
+        make_rust_extension("rustapi_module.pyclass_iter"),
     ],
     install_requires=install_requires,
     tests_require=tests_require,

examples/rustapi_module/src/lib.rs

@@ -3,4 +3,5 @@ pub mod datetime;
 pub mod dict_iter;
 pub mod objstore;
 pub mod othermod;
+pub mod pyclass_iter;
 pub mod subclassing;

examples/rustapi_module/src/pyclass_iter.rs

@@ -0,0 +1,34 @@
+use pyo3::class::iter::{IterNextOutput, PyIterProtocol};
+use pyo3::prelude::*;
+
+/// This is for demonstrating how to return a value from __next__
+#[pyclass]
+struct PyClassIter {
+    count: usize,
+}
+
+#[pymethods]
+impl PyClassIter {
+    #[new]
+    pub fn new() -> Self {
+        PyClassIter { count: 0 }
+    }
+}
+
+#[pyproto]
+impl PyIterProtocol for PyClassIter {
+    fn __next__(mut slf: PyRefMut<Self>) -> IterNextOutput<usize, &'static str> {
+        if slf.count < 5 {
+            slf.count += 1;
+            IterNextOutput::Yield(slf.count)
+        } else {
+            IterNextOutput::Return("Ended")
+        }
+    }
+}
+
+#[pymodule]
+pub fn pyclass_iter(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<PyClassIter>()?;
+    Ok(())
+}

examples/rustapi_module/tests/test_pyclass_iter.py

@@ -0,0 +1,15 @@
+import pytest
+from rustapi_module import pyclass_iter
+
+
+def test_iter():
+    i = pyclass_iter.PyClassIter()
+    assert next(i) == 1
+    assert next(i) == 2
+    assert next(i) == 3
+    assert next(i) == 4
+    assert next(i) == 5
+
+    with pytest.raises(StopIteration) as excinfo:
+        next(i)
+    assert excinfo.value.value == "Ended"

guide/src/class.md

@@ -803,11 +803,10 @@ It includes two methods `__iter__` and `__next__`:
   * `fn __iter__(slf: PyRefMut<Self>) -> PyResult<impl IntoPy<PyObject>>`
   * `fn __next__(slf: PyRefMut<Self>) -> PyResult<Option<impl IntoPy<PyObject>>>`
 
-Returning `Ok(None)` from `__next__` indicates that that there are no further items.
+Returning `None` from `__next__` indicates that that there are no further items.
 These two methods can be take either `PyRef<Self>` or `PyRefMut<Self>` as their
 first argument, so that mutable borrow can be avoided if needed.
 
-
 Example:
 
 ```rust
@@ -886,6 +885,14 @@ impl PyIterProtocol for Container {
 For more details on Python's iteration protocols, check out [the "Iterator Types" section of the library
 documentation](https://docs.python.org/3/library/stdtypes.html#iterator-types).
 
+#### Returning a value from iteration
+
+This guide has so far shown how to use `Option<T>` to implement yielding values during iteration.
+In Python a generator can also return a value. To express this in Rust, PyO3 provides the
+[`IterNextOutput`](https://docs.rs/pyo3/latest/pyo3/class/iter/enum.IterNextOutput.html) enum to
+both `Yield` values and `Return` a final value - see its docs for further details and an example.
+
+
 ## How methods are implemented
 
 Users should be able to define a `#[pyclass]` with or without `#[pymethods]`, while PyO3 needs a

src/class/iter.rs

@@ -11,6 +11,39 @@ use crate::{ffi, IntoPy, IntoPyPointer, PyClass, PyObject, Python};
 ///
 /// Check [CPython doc](https://docs.python.org/3/c-api/typeobj.html#c.PyTypeObject.tp_iter)
 /// for more.
+///
+/// # Example
+/// The following example shows how to implement a simple Python iterator in Rust which yields
+/// the integers 1 to 5, before raising `StopIteration("Ended")`.
+///
+/// ```rust
+/// use pyo3::prelude::*;
+/// use pyo3::PyIterProtocol;
+/// use pyo3::class::iter::IterNextOutput;
+///
+/// #[pyclass]
+/// struct Iter {
+///     count: usize
+/// }
+///
+/// #[pyproto]
+/// impl PyIterProtocol for Iter {
+///     fn __next__(mut slf: PyRefMut<Self>) -> IterNextOutput<usize, &'static str> {
+///         if slf.count < 5 {
+///             slf.count += 1;
+///             IterNextOutput::Yield(slf.count)
+///         } else {
+///             IterNextOutput::Return("Ended")
+///         }
+///     }
+/// }
+///
+/// # let gil = Python::acquire_gil();
+/// # let py = gil.python();
+/// # let inst = Py::new(py, Iter { count: 0 }).unwrap();
+/// # pyo3::py_run!(py, inst, "assert next(inst) == 1");
+/// # // test of StopIteration is done in examples/rustapi_module/pyclass_iter.rs
+/// ```
 #[allow(unused_variables)]
 pub trait PyIterProtocol<'p>: PyClass {
     fn __iter__(slf: Self::Receiver) -> Self::Result
@@ -35,7 +68,7 @@ pub trait PyIterIterProtocol<'p>: PyIterProtocol<'p> {
 
 pub trait PyIterNextProtocol<'p>: PyIterProtocol<'p> {
     type Receiver: TryFromPyCell<'p, Self>;
-    type Result: IntoPyCallbackOutput<IterNextOutput>;
+    type Result: IntoPyCallbackOutput<PyIterNextOutput>;
 }
 
 #[derive(Default)]
@@ -64,22 +97,47 @@ impl PyIterMethods {
     }
 }
 
-pub struct IterNextOutput(Option<PyObject>);
+/// Output of `__next__` which can either `yield` the next value in the iteration, or
+/// `return` a value to raise `StopIteration` in Python.
+///
+/// See [`PyIterProtocol`](trait.PyIterProtocol.html) for an example.
+pub enum IterNextOutput<T, U> {
+    Yield(T),
+    Return(U),
+}
 
-impl IntoPyCallbackOutput<*mut ffi::PyObject> for IterNextOutput {
+pub type PyIterNextOutput = IterNextOutput<PyObject, PyObject>;
+
+impl IntoPyCallbackOutput<*mut ffi::PyObject> for PyIterNextOutput {
     fn convert(self, _py: Python) -> PyResult<*mut ffi::PyObject> {
-        match self.0 {
+        match self {
-            Some(o) => Ok(o.into_ptr()),
+            IterNextOutput::Yield(o) => Ok(o.into_ptr()),
-            None => Err(crate::exceptions::StopIteration::py_err(())),
+            IterNextOutput::Return(opt) => Err(crate::exceptions::StopIteration::py_err((opt,))),
         }
     }
 }
 
-impl<T> IntoPyCallbackOutput<IterNextOutput> for Option<T>
+impl<T, U> IntoPyCallbackOutput<PyIterNextOutput> for IterNextOutput<T, U>
+where
+    T: IntoPy<PyObject>,
+    U: IntoPy<PyObject>,
+{
+    fn convert(self, py: Python) -> PyResult<PyIterNextOutput> {
+        match self {
+            IterNextOutput::Yield(o) => Ok(IterNextOutput::Yield(o.into_py(py))),
+            IterNextOutput::Return(o) => Ok(IterNextOutput::Return(o.into_py(py))),
+        }
+    }
+}
+
+impl<T> IntoPyCallbackOutput<PyIterNextOutput> for Option<T>
 where
     T: IntoPy<PyObject>,
 {
-    fn convert(self, py: Python) -> PyResult<IterNextOutput> {
+    fn convert(self, py: Python) -> PyResult<PyIterNextOutput> {
-        Ok(IterNextOutput(self.map(|o| o.into_py(py))))
+        match self {
+            Some(o) => Ok(PyIterNextOutput::Yield(o.into_py(py))),
+            None => Ok(PyIterNextOutput::Return(py.None())),
+        }
     }
 }

src/class/pyasync.rs

@@ -71,7 +71,7 @@ pub trait PyAsyncAiterProtocol<'p>: PyAsyncProtocol<'p> {
 
 pub trait PyAsyncAnextProtocol<'p>: PyAsyncProtocol<'p> {
     type Receiver: TryFromPyCell<'p, Self>;
-    type Result: IntoPyCallbackOutput<IterANextOutput>;
+    type Result: IntoPyCallbackOutput<PyIterANextOutput>;
 }
 
 pub trait PyAsyncAenterProtocol<'p>: PyAsyncProtocol<'p> {
@@ -107,23 +107,46 @@ impl ffi::PyAsyncMethods {
     }
 }
 
-pub struct IterANextOutput(Option<PyObject>);
+pub enum IterANextOutput<T, U> {
+    Yield(T),
+    Return(U),
+}
 
-impl IntoPyCallbackOutput<*mut ffi::PyObject> for IterANextOutput {
+pub type PyIterANextOutput = IterANextOutput<PyObject, PyObject>;
+
+impl IntoPyCallbackOutput<*mut ffi::PyObject> for PyIterANextOutput {
     fn convert(self, _py: Python) -> PyResult<*mut ffi::PyObject> {
-        match self.0 {
+        match self {
-            Some(o) => Ok(o.into_ptr()),
+            IterANextOutput::Yield(o) => Ok(o.into_ptr()),
-            None => Err(crate::exceptions::StopAsyncIteration::py_err(())),
+            IterANextOutput::Return(opt) => {
+                Err(crate::exceptions::StopAsyncIteration::py_err((opt,)))
+            }
         }
     }
 }
 
-impl<T> IntoPyCallbackOutput<IterANextOutput> for Option<T>
+impl<T, U> IntoPyCallbackOutput<PyIterANextOutput> for IterANextOutput<T, U>
+where
+    T: IntoPy<PyObject>,
+    U: IntoPy<PyObject>,
+{
+    fn convert(self, py: Python) -> PyResult<PyIterANextOutput> {
+        match self {
+            IterANextOutput::Yield(o) => Ok(IterANextOutput::Yield(o.into_py(py))),
+            IterANextOutput::Return(o) => Ok(IterANextOutput::Return(o.into_py(py))),
+        }
+    }
+}
+
+impl<T> IntoPyCallbackOutput<PyIterANextOutput> for Option<T>
 where
     T: IntoPy<PyObject>,
 {
-    fn convert(self, py: Python) -> PyResult<IterANextOutput> {
+    fn convert(self, py: Python) -> PyResult<PyIterANextOutput> {
-        Ok(IterANextOutput(self.map(|o| o.into_py(py))))
+        match self {
+            Some(o) => Ok(PyIterANextOutput::Yield(o.into_py(py))),
+            None => Ok(PyIterANextOutput::Return(py.None())),
+        }
     }
 }