pyo3/types/
iterator.rs

1use crate::ffi_ptr_ext::FfiPtrExt;
2use crate::py_result_ext::PyResultExt;
3use crate::sync::PyOnceLock;
4#[cfg(Py_LIMITED_API)]
5use crate::types::PyAnyMethods;
6use crate::types::{PyType, PyTypeMethods};
7use crate::{ffi, Bound, Py, PyAny, PyErr, PyResult};
8
9/// A Python iterator object.
10///
11/// Values of this type are accessed via PyO3's smart pointers, e.g. as
12/// [`Py<PyIterator>`][crate::Py] or [`Bound<'py, PyIterator>`][Bound].
13///
14/// # Examples
15///
16/// ```rust
17/// use pyo3::prelude::*;
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 = core::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    fn next(&mut self) -> Option<Self::Item> {
107        let py = self.py();
108        let mut item = core::ptr::null_mut();
109
110        // SAFETY: `self` is a valid iterator object, `item` is a valid pointer to receive the next item
111        match unsafe { ffi::compat::PyIter_NextItem(self.as_ptr(), &mut item) } {
112            core::ffi::c_int::MIN..=-1 => Some(Err(PyErr::fetch(py))),
113            0 => None,
114            // SAFETY: `item` is guaranteed to be a non-null strong reference
115            1..=core::ffi::c_int::MAX => Some(Ok(unsafe { item.assume_owned_unchecked(py) })),
116        }
117    }
118
119    fn size_hint(&self) -> (usize, Option<usize>) {
120        match length_hint(self) {
121            Ok(hint) => (hint, None),
122            Err(e) => {
123                e.write_unraisable(self.py(), Some(self));
124                (0, None)
125            }
126        }
127    }
128}
129
130#[cfg(not(Py_LIMITED_API))]
131fn length_hint(iter: &Bound<'_, PyIterator>) -> PyResult<usize> {
132    // SAFETY: `iter` is a valid iterator object
133    let hint = unsafe { ffi::PyObject_LengthHint(iter.as_ptr(), 0) };
134    if hint < 0 {
135        Err(PyErr::fetch(iter.py()))
136    } else {
137        Ok(hint as usize)
138    }
139}
140
141/// On the limited API, we cannot use `PyObject_LengthHint`, so we fall back to calling
142/// `operator.length_hint()`, which is documented equivalent to calling `PyObject_LengthHint`.
143#[cfg(Py_LIMITED_API)]
144fn length_hint(iter: &Bound<'_, PyIterator>) -> PyResult<usize> {
145    static LENGTH_HINT: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
146    let length_hint = LENGTH_HINT.import(iter.py(), "operator", "length_hint")?;
147    length_hint.call1((iter, 0))?.extract()
148}
149
150impl<'py> IntoIterator for &Bound<'py, PyIterator> {
151    type Item = PyResult<Bound<'py, PyAny>>;
152    type IntoIter = Bound<'py, PyIterator>;
153
154    fn into_iter(self) -> Self::IntoIter {
155        self.clone()
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::PyIterator;
162    #[cfg(all(not(PyPy), Py_3_10))]
163    use super::PySendResult;
164    use crate::exceptions::PyTypeError;
165    #[cfg(all(not(PyPy), Py_3_10))]
166    use crate::types::PyNone;
167    use crate::types::{PyAnyMethods, PyDict, PyList, PyListMethods};
168    #[cfg(feature = "macros")]
169    use crate::PyErr;
170    use crate::{IntoPyObject, PyTypeInfo, Python};
171
172    #[test]
173    fn vec_iter() {
174        Python::attach(|py| {
175            let inst = vec![10, 20].into_pyobject(py).unwrap();
176            let mut it = inst.try_iter().unwrap();
177            assert_eq!(
178                10_i32,
179                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
180            );
181            assert_eq!(
182                20_i32,
183                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
184            );
185            assert!(it.next().is_none());
186        });
187    }
188
189    #[test]
190    fn iter_refcnt() {
191        let (obj, count) = Python::attach(|py| {
192            let obj = vec![10, 20].into_pyobject(py).unwrap();
193            let count = obj._get_refcnt();
194            (obj.unbind(), count)
195        });
196
197        Python::attach(|py| {
198            let inst = obj.bind(py);
199            let mut it = inst.try_iter().unwrap();
200
201            assert_eq!(
202                10_i32,
203                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
204            );
205        });
206
207        Python::attach(move |py| {
208            assert_eq!(count, obj._get_refcnt(py));
209        });
210    }
211
212    #[test]
213    fn iter_item_refcnt() {
214        Python::attach(|py| {
215            let count;
216            let obj = py.eval(c"object()", None, None).unwrap();
217            let list = {
218                let list = PyList::empty(py);
219                list.append(10).unwrap();
220                list.append(&obj).unwrap();
221                count = obj._get_refcnt();
222                list
223            };
224
225            {
226                let mut it = list.iter();
227
228                assert_eq!(10_i32, it.next().unwrap().extract::<'_, i32>().unwrap());
229                assert!(it.next().unwrap().is(&obj));
230                assert!(it.next().is_none());
231            }
232            assert_eq!(count, obj._get_refcnt());
233        });
234    }
235
236    #[test]
237    fn fibonacci_generator() {
238        let fibonacci_generator = cr#"
239def fibonacci(target):
240    a = 1
241    b = 1
242    for _ in range(target):
243        yield a
244        a, b = b, a + b
245"#;
246
247        Python::attach(|py| {
248            let context = PyDict::new(py);
249            py.run(fibonacci_generator, None, Some(&context)).unwrap();
250
251            let generator = py.eval(c"fibonacci(5)", None, Some(&context)).unwrap();
252            for (actual, expected) in generator.try_iter().unwrap().zip(&[1, 1, 2, 3, 5]) {
253                let actual = actual.unwrap().extract::<usize>().unwrap();
254                assert_eq!(actual, *expected)
255            }
256        });
257    }
258
259    #[test]
260    #[cfg(all(not(PyPy), Py_3_10))]
261    fn send_generator() {
262        let generator = cr#"
263def gen():
264    value = None
265    while(True):
266        value = yield value
267        if value is None:
268            return
269"#;
270
271        Python::attach(|py| {
272            let context = PyDict::new(py);
273            py.run(generator, None, Some(&context)).unwrap();
274
275            let generator = py.eval(c"gen()", None, Some(&context)).unwrap();
276
277            let one = 1i32.into_pyobject(py).unwrap();
278            assert!(matches!(
279                generator.try_iter().unwrap().send(&PyNone::get(py)).unwrap(),
280                PySendResult::Next(value) if value.is_none()
281            ));
282            assert!(matches!(
283                generator.try_iter().unwrap().send(&one).unwrap(),
284                PySendResult::Next(value) if value.is(&one)
285            ));
286            assert!(matches!(
287                generator.try_iter().unwrap().send(&PyNone::get(py)).unwrap(),
288                PySendResult::Return(value) if value.is_none()
289            ));
290        });
291    }
292
293    #[test]
294    fn fibonacci_generator_bound() {
295        use crate::types::any::PyAnyMethods;
296        use crate::Bound;
297
298        let fibonacci_generator = cr#"
299def fibonacci(target):
300    a = 1
301    b = 1
302    for _ in range(target):
303        yield a
304        a, b = b, a + b
305"#;
306
307        Python::attach(|py| {
308            let context = PyDict::new(py);
309            py.run(fibonacci_generator, None, Some(&context)).unwrap();
310
311            let generator: Bound<'_, PyIterator> = py
312                .eval(c"fibonacci(5)", None, Some(&context))
313                .unwrap()
314                .cast_into()
315                .unwrap();
316            let mut items = vec![];
317            for actual in &generator {
318                let actual = actual.unwrap().extract::<usize>().unwrap();
319                items.push(actual);
320            }
321            assert_eq!(items, [1, 1, 2, 3, 5]);
322        });
323    }
324
325    #[test]
326    fn int_not_iterable() {
327        Python::attach(|py| {
328            let x = 5i32.into_pyobject(py).unwrap();
329            let err = PyIterator::from_object(&x).unwrap_err();
330
331            assert!(err.is_instance_of::<PyTypeError>(py));
332        });
333    }
334
335    #[test]
336    #[cfg(feature = "macros")]
337    fn python_class_not_iterator() {
338        use crate::PyErr;
339
340        #[crate::pyclass(crate = "crate")]
341        struct Downcaster {
342            failed: Option<PyErr>,
343        }
344
345        #[crate::pymethods(crate = "crate")]
346        impl Downcaster {
347            fn downcast_iterator(&mut self, obj: &crate::Bound<'_, crate::PyAny>) {
348                self.failed = Some(obj.cast::<PyIterator>().unwrap_err().into());
349            }
350        }
351
352        // Regression test for 2913
353        Python::attach(|py| {
354            let downcaster = crate::Py::new(py, Downcaster { failed: None }).unwrap();
355            crate::py_run!(
356                py,
357                downcaster,
358                r#"
359                    from collections.abc import Sequence
360
361                    class MySequence(Sequence):
362                        def __init__(self):
363                            self._data = [1, 2, 3]
364
365                        def __getitem__(self, index):
366                            return self._data[index]
367
368                        def __len__(self):
369                            return len(self._data)
370
371                    downcaster.downcast_iterator(MySequence())
372                "#
373            );
374
375            assert_eq!(
376                downcaster.borrow_mut(py).failed.take().unwrap().to_string(),
377                "TypeError: 'MySequence' object is not an instance of 'Iterator'"
378            );
379        });
380    }
381
382    #[test]
383    #[cfg(feature = "macros")]
384    fn python_class_iterator() {
385        #[crate::pyfunction(crate = "crate")]
386        fn assert_iterator(obj: &crate::Bound<'_, crate::PyAny>) {
387            assert!(obj.cast::<PyIterator>().is_ok())
388        }
389
390        // Regression test for 2913
391        Python::attach(|py| {
392            let assert_iterator = crate::wrap_pyfunction!(assert_iterator, py).unwrap();
393            crate::py_run!(
394                py,
395                assert_iterator,
396                r#"
397                    class MyIter:
398                        def __next__(self):
399                            raise StopIteration
400
401                    assert_iterator(MyIter())
402                "#
403            );
404        });
405    }
406
407    #[test]
408    fn length_hint_becomes_size_hint_lower_bound() {
409        Python::attach(|py| {
410            let list = py.eval(c"[1, 2, 3]", None, None).unwrap();
411            let iter = list.try_iter().unwrap();
412            let hint = iter.size_hint();
413            assert_eq!(hint, (3, None));
414        });
415    }
416
417    #[test]
418    #[cfg(feature = "macros")]
419    fn length_hint_error() {
420        #[crate::pyfunction(crate = "crate")]
421        fn test_size_hint(obj: &crate::Bound<'_, crate::PyAny>, should_error: bool) {
422            let iter = obj.cast::<PyIterator>().unwrap();
423            crate::test_utils::UnraisableCapture::enter(obj.py(), |capture| {
424                assert_eq!((0, None), iter.size_hint());
425                assert_eq!(should_error, capture.take_capture().is_some());
426            });
427            assert!(PyErr::take(obj.py()).is_none());
428        }
429
430        Python::attach(|py| {
431            let test_size_hint = crate::wrap_pyfunction!(test_size_hint, py).unwrap();
432            crate::py_run!(
433                py,
434                test_size_hint,
435                r#"
436                    class NoHintIter:
437                        def __next__(self):
438                            raise StopIteration
439
440                        def __length_hint__(self):
441                            return NotImplemented
442
443                    class ErrorHintIter:
444                        def __next__(self):
445                            raise StopIteration
446
447                        def __length_hint__(self):
448                            raise ValueError("bad hint impl")
449
450                    test_size_hint(NoHintIter(), False)
451                    test_size_hint(ErrorHintIter(), True)
452                "#
453            );
454        });
455    }
456
457    #[test]
458    fn test_type_object() {
459        Python::attach(|py| {
460            let abc = PyIterator::type_object(py);
461            let iter = py.eval(c"iter(())", None, None).unwrap();
462            assert!(iter.is_instance(&abc).unwrap());
463        })
464    }
465}
pyo3/types/iterator.rs

pyo3/types/
iterator.rs
