pyo3/types/
iterator.rs

1use crate::ffi_ptr_ext::FfiPtrExt;
2use crate::instance::Borrowed;
3use crate::py_result_ext::PyResultExt;
4use crate::{ffi, Bound, PyAny, PyErr, PyResult, PyTypeCheck};
5
6/// A Python iterator object.
7///
8/// Values of this type are accessed via PyO3's smart pointers, e.g. as
9/// [`Py<PyIterator>`][crate::Py] or [`Bound<'py, PyIterator>`][Bound].
10///
11/// # Examples
12///
13/// ```rust
14/// use pyo3::prelude::*;
15/// use pyo3::ffi::c_str;
16///
17/// # fn main() -> PyResult<()> {
18/// Python::with_gil(|py| -> PyResult<()> {
19///     let list = py.eval(c_str!("iter([1, 2, 3, 4])"), None, None)?;
20///     let numbers: PyResult<Vec<usize>> = list
21///         .try_iter()?
22///         .map(|i| i.and_then(|i|i.extract::<usize>()))
23///         .collect();
24///     let sum: usize = numbers?.iter().sum();
25///     assert_eq!(sum, 10);
26///     Ok(())
27/// })
28/// # }
29/// ```
30#[repr(transparent)]
31pub struct PyIterator(PyAny);
32pyobject_native_type_named!(PyIterator);
33
34impl PyIterator {
35    /// Builds an iterator for an iterable Python object; the equivalent of calling `iter(obj)` in Python.
36    ///
37    /// Usually it is more convenient to write [`obj.try_iter()`][crate::types::any::PyAnyMethods::try_iter],
38    /// which is a more concise way of calling this function.
39    pub fn from_object<'py>(obj: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyIterator>> {
40        unsafe {
41            ffi::PyObject_GetIter(obj.as_ptr())
42                .assume_owned_or_err(obj.py())
43                .downcast_into_unchecked()
44        }
45    }
46}
47
48#[derive(Debug)]
49#[cfg(all(not(PyPy), Py_3_10))]
50pub enum PySendResult<'py> {
51    Next(Bound<'py, PyAny>),
52    Return(Bound<'py, PyAny>),
53}
54
55#[cfg(all(not(PyPy), Py_3_10))]
56impl<'py> Bound<'py, PyIterator> {
57    /// Sends a value into a python generator. This is the equivalent of calling `generator.send(value)` in Python.
58    /// This resumes the generator and continues its execution until the next `yield` or `return` statement.
59    /// If the generator exits without returning a value, this function returns a `StopException`.
60    /// The first call to `send` must be made with `None` as the argument to start the generator, failing to do so will raise a `TypeError`.
61    #[inline]
62    pub fn send(&self, value: &Bound<'py, PyAny>) -> PyResult<PySendResult<'py>> {
63        let py = self.py();
64        let mut result = std::ptr::null_mut();
65        match unsafe { ffi::PyIter_Send(self.as_ptr(), value.as_ptr(), &mut result) } {
66            ffi::PySendResult::PYGEN_ERROR => Err(PyErr::fetch(py)),
67            ffi::PySendResult::PYGEN_RETURN => Ok(PySendResult::Return(unsafe {
68                result.assume_owned_unchecked(py)
69            })),
70            ffi::PySendResult::PYGEN_NEXT => Ok(PySendResult::Next(unsafe {
71                result.assume_owned_unchecked(py)
72            })),
73        }
74    }
75}
76
77impl<'py> Iterator for Bound<'py, PyIterator> {
78    type Item = PyResult<Bound<'py, PyAny>>;
79
80    /// Retrieves the next item from an iterator.
81    ///
82    /// Returns `None` when the iterator is exhausted.
83    /// If an exception occurs, returns `Some(Err(..))`.
84    /// Further `next()` calls after an exception occurs are likely
85    /// to repeatedly result in the same exception.
86    #[inline]
87    fn next(&mut self) -> Option<Self::Item> {
88        Borrowed::from(&*self).next()
89    }
90
91    #[cfg(not(Py_LIMITED_API))]
92    fn size_hint(&self) -> (usize, Option<usize>) {
93        let hint = unsafe { ffi::PyObject_LengthHint(self.as_ptr(), 0) };
94        (hint.max(0) as usize, None)
95    }
96}
97
98impl<'py> Borrowed<'_, 'py, PyIterator> {
99    // TODO: this method is on Borrowed so that &'py PyIterator can use this; once that
100    // implementation is deleted this method should be moved to the `Bound<'py, PyIterator> impl
101    fn next(self) -> Option<PyResult<Bound<'py, PyAny>>> {
102        let py = self.py();
103
104        match unsafe { ffi::PyIter_Next(self.as_ptr()).assume_owned_or_opt(py) } {
105            Some(obj) => Some(Ok(obj)),
106            None => PyErr::take(py).map(Err),
107        }
108    }
109}
110
111impl<'py> IntoIterator for &Bound<'py, PyIterator> {
112    type Item = PyResult<Bound<'py, PyAny>>;
113    type IntoIter = Bound<'py, PyIterator>;
114
115    fn into_iter(self) -> Self::IntoIter {
116        self.clone()
117    }
118}
119
120impl PyTypeCheck for PyIterator {
121    const NAME: &'static str = "Iterator";
122
123    fn type_check(object: &Bound<'_, PyAny>) -> bool {
124        unsafe { ffi::PyIter_Check(object.as_ptr()) != 0 }
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::PyIterator;
131    #[cfg(all(not(PyPy), Py_3_10))]
132    use super::PySendResult;
133    use crate::exceptions::PyTypeError;
134    #[cfg(all(not(PyPy), Py_3_10))]
135    use crate::types::PyNone;
136    use crate::types::{PyAnyMethods, PyDict, PyList, PyListMethods};
137    use crate::{ffi, IntoPyObject, Python};
138
139    #[test]
140    fn vec_iter() {
141        Python::with_gil(|py| {
142            let inst = vec![10, 20].into_pyobject(py).unwrap();
143            let mut it = inst.try_iter().unwrap();
144            assert_eq!(
145                10_i32,
146                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
147            );
148            assert_eq!(
149                20_i32,
150                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
151            );
152            assert!(it.next().is_none());
153        });
154    }
155
156    #[test]
157    fn iter_refcnt() {
158        let (obj, count) = Python::with_gil(|py| {
159            let obj = vec![10, 20].into_pyobject(py).unwrap();
160            let count = obj.get_refcnt();
161            (obj.unbind(), count)
162        });
163
164        Python::with_gil(|py| {
165            let inst = obj.bind(py);
166            let mut it = inst.try_iter().unwrap();
167
168            assert_eq!(
169                10_i32,
170                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
171            );
172        });
173
174        Python::with_gil(move |py| {
175            assert_eq!(count, obj.get_refcnt(py));
176        });
177    }
178
179    #[test]
180    fn iter_item_refcnt() {
181        Python::with_gil(|py| {
182            let count;
183            let obj = py.eval(ffi::c_str!("object()"), None, None).unwrap();
184            let list = {
185                let list = PyList::empty(py);
186                list.append(10).unwrap();
187                list.append(&obj).unwrap();
188                count = obj.get_refcnt();
189                list
190            };
191
192            {
193                let mut it = list.iter();
194
195                assert_eq!(10_i32, it.next().unwrap().extract::<'_, i32>().unwrap());
196                assert!(it.next().unwrap().is(&obj));
197                assert!(it.next().is_none());
198            }
199            assert_eq!(count, obj.get_refcnt());
200        });
201    }
202
203    #[test]
204    fn fibonacci_generator() {
205        let fibonacci_generator = ffi::c_str!(
206            r#"
207def fibonacci(target):
208    a = 1
209    b = 1
210    for _ in range(target):
211        yield a
212        a, b = b, a + b
213"#
214        );
215
216        Python::with_gil(|py| {
217            let context = PyDict::new(py);
218            py.run(fibonacci_generator, None, Some(&context)).unwrap();
219
220            let generator = py
221                .eval(ffi::c_str!("fibonacci(5)"), None, Some(&context))
222                .unwrap();
223            for (actual, expected) in generator.try_iter().unwrap().zip(&[1, 1, 2, 3, 5]) {
224                let actual = actual.unwrap().extract::<usize>().unwrap();
225                assert_eq!(actual, *expected)
226            }
227        });
228    }
229
230    #[test]
231    #[cfg(all(not(PyPy), Py_3_10))]
232    fn send_generator() {
233        let generator = ffi::c_str!(
234            r#"
235def gen():
236    value = None
237    while(True):
238        value = yield value
239        if value is None:
240            return
241"#
242        );
243
244        Python::with_gil(|py| {
245            let context = PyDict::new(py);
246            py.run(generator, None, Some(&context)).unwrap();
247
248            let generator = py.eval(ffi::c_str!("gen()"), None, Some(&context)).unwrap();
249
250            let one = 1i32.into_pyobject(py).unwrap();
251            assert!(matches!(
252                generator.try_iter().unwrap().send(&PyNone::get(py)).unwrap(),
253                PySendResult::Next(value) if value.is_none()
254            ));
255            assert!(matches!(
256                generator.try_iter().unwrap().send(&one).unwrap(),
257                PySendResult::Next(value) if value.is(&one)
258            ));
259            assert!(matches!(
260                generator.try_iter().unwrap().send(&PyNone::get(py)).unwrap(),
261                PySendResult::Return(value) if value.is_none()
262            ));
263        });
264    }
265
266    #[test]
267    fn fibonacci_generator_bound() {
268        use crate::types::any::PyAnyMethods;
269        use crate::Bound;
270
271        let fibonacci_generator = ffi::c_str!(
272            r#"
273def fibonacci(target):
274    a = 1
275    b = 1
276    for _ in range(target):
277        yield a
278        a, b = b, a + b
279"#
280        );
281
282        Python::with_gil(|py| {
283            let context = PyDict::new(py);
284            py.run(fibonacci_generator, None, Some(&context)).unwrap();
285
286            let generator: Bound<'_, PyIterator> = py
287                .eval(ffi::c_str!("fibonacci(5)"), None, Some(&context))
288                .unwrap()
289                .downcast_into()
290                .unwrap();
291            let mut items = vec![];
292            for actual in &generator {
293                let actual = actual.unwrap().extract::<usize>().unwrap();
294                items.push(actual);
295            }
296            assert_eq!(items, [1, 1, 2, 3, 5]);
297        });
298    }
299
300    #[test]
301    fn int_not_iterable() {
302        Python::with_gil(|py| {
303            let x = 5i32.into_pyobject(py).unwrap();
304            let err = PyIterator::from_object(&x).unwrap_err();
305
306            assert!(err.is_instance_of::<PyTypeError>(py));
307        });
308    }
309
310    #[test]
311    #[cfg(feature = "macros")]
312    fn python_class_not_iterator() {
313        use crate::PyErr;
314
315        #[crate::pyclass(crate = "crate")]
316        struct Downcaster {
317            failed: Option<PyErr>,
318        }
319
320        #[crate::pymethods(crate = "crate")]
321        impl Downcaster {
322            fn downcast_iterator(&mut self, obj: &crate::Bound<'_, crate::PyAny>) {
323                self.failed = Some(obj.downcast::<PyIterator>().unwrap_err().into());
324            }
325        }
326
327        // Regression test for 2913
328        Python::with_gil(|py| {
329            let downcaster = crate::Py::new(py, Downcaster { failed: None }).unwrap();
330            crate::py_run!(
331                py,
332                downcaster,
333                r#"
334                    from collections.abc import Sequence
335
336                    class MySequence(Sequence):
337                        def __init__(self):
338                            self._data = [1, 2, 3]
339
340                        def __getitem__(self, index):
341                            return self._data[index]
342
343                        def __len__(self):
344                            return len(self._data)
345
346                    downcaster.downcast_iterator(MySequence())
347                "#
348            );
349
350            assert_eq!(
351                downcaster.borrow_mut(py).failed.take().unwrap().to_string(),
352                "TypeError: 'MySequence' object cannot be converted to 'Iterator'"
353            );
354        });
355    }
356
357    #[test]
358    #[cfg(feature = "macros")]
359    fn python_class_iterator() {
360        #[crate::pyfunction(crate = "crate")]
361        fn assert_iterator(obj: &crate::Bound<'_, crate::PyAny>) {
362            assert!(obj.downcast::<PyIterator>().is_ok())
363        }
364
365        // Regression test for 2913
366        Python::with_gil(|py| {
367            let assert_iterator = crate::wrap_pyfunction!(assert_iterator, py).unwrap();
368            crate::py_run!(
369                py,
370                assert_iterator,
371                r#"
372                    class MyIter:
373                        def __next__(self):
374                            raise StopIteration
375
376                    assert_iterator(MyIter())
377                "#
378            );
379        });
380    }
381
382    #[test]
383    #[cfg(not(Py_LIMITED_API))]
384    fn length_hint_becomes_size_hint_lower_bound() {
385        Python::with_gil(|py| {
386            let list = py.eval(ffi::c_str!("[1, 2, 3]"), None, None).unwrap();
387            let iter = list.try_iter().unwrap();
388            let hint = iter.size_hint();
389            assert_eq!(hint, (3, None));
390        });
391    }
392}
pyo3/types/iterator.rs

pyo3/types/
iterator.rs
