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

pyo3/types/
iterator.rs
