pyo3/types/
iterator.rs

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

pyo3/types/
iterator.rs
