Add ability to return from __next__ / __anext__

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 @@ def make_rust_extension(module_name):
         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 {
+        if slf.count < 5 {
+            slf.count += 1;
+            IterNextOutput::Yield(slf.count.into_py(slf.py()))
+        } else {
+            IterNextOutput::Return(Some("Ended".into_py(slf.py())))
+        }
+    }
+}
+
+#[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,14 @@
+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"

src/class/iter.rs

@@ -64,22 +64,37 @@ impl PyIterMethods {
     }
 }
 
-pub struct IterNextOutput(Option<PyObject>);
+pub enum IterNextOutput {
+    Yield(PyObject),
+    Return(Option<PyObject>),
+}
 
 impl IntoPyCallbackOutput<*mut ffi::PyObject> for IterNextOutput {
     fn convert(self, _py: Python) -> PyResult<*mut ffi::PyObject> {
-        match self.0 {
-            Some(o) => Ok(o.into_ptr()),
-            None => Err(crate::exceptions::StopIteration::py_err(())),
+        match self {
+            IterNextOutput::Yield(o) => Ok(o.into_ptr()),
+            IterNextOutput::Return(opt) => match opt {
+                Some(o) => Err(crate::exceptions::StopIteration::py_err((o,))),
+                None => Err(crate::exceptions::StopIteration::py_err(())),
+            },
         }
     }
 }
 
+impl IntoPyCallbackOutput<IterNextOutput> for IterNextOutput {
+    fn convert(self, _py: Python) -> PyResult<IterNextOutput> {
+        Ok(self)
+    }
+}
+
 impl<T> IntoPyCallbackOutput<IterNextOutput> for Option<T>
 where
     T: IntoPy<PyObject>,
 {
     fn convert(self, py: Python) -> PyResult<IterNextOutput> {
-        Ok(IterNextOutput(self.map(|o| o.into_py(py))))
+        match self {
+            Some(o) => Ok(IterNextOutput::Yield(o.into_py(py))),
+            None => Ok(IterNextOutput::Return(None)),
+        }
     }
 }

src/class/pyasync.rs

@@ -107,23 +107,38 @@ impl ffi::PyAsyncMethods {
     }
 }
 
-pub struct IterANextOutput(Option<PyObject>);
+pub enum IterANextOutput {
+    Yield(PyObject),
+    Return(Option<PyObject>),
+}
 
 impl IntoPyCallbackOutput<*mut ffi::PyObject> for IterANextOutput {
     fn convert(self, _py: Python) -> PyResult<*mut ffi::PyObject> {
-        match self.0 {
-            Some(o) => Ok(o.into_ptr()),
-            None => Err(crate::exceptions::StopAsyncIteration::py_err(())),
+        match self {
+            IterANextOutput::Yield(o) => Ok(o.into_ptr()),
+            IterANextOutput::Return(opt) => match opt {
+                Some(o) => Err(crate::exceptions::StopAsyncIteration::py_err((o,))),
+                None => Err(crate::exceptions::StopAsyncIteration::py_err(())),
+            },
         }
     }
 }
 
+impl IntoPyCallbackOutput<IterANextOutput> for IterANextOutput {
+    fn convert(self, _py: Python) -> PyResult<IterANextOutput> {
+        Ok(self)
+    }
+}
+
 impl<T> IntoPyCallbackOutput<IterANextOutput> for Option<T>
 where
     T: IntoPy<PyObject>,
 {
     fn convert(self, py: Python) -> PyResult<IterANextOutput> {
-        Ok(IterANextOutput(self.map(|o| o.into_py(py))))
+        match self {
+            Some(o) => Ok(IterANextOutput::Yield(o.into_py(py))),
+            None => Ok(IterANextOutput::Return(None)),
+        }
     }
 }