pyo3/types/
iterator.rs

Help
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
55impl<'py> Iterator for Bound<'py, PyIterator> {
56    type Item = PyResult<Bound<'py, PyAny>>;
57
58    /// Retrieves the next item from an iterator.
59    ///
60    /// Returns `None` when the iterator is exhausted.
61    /// If an exception occurs, returns `Some(Err(..))`.
62    /// Further `next()` calls after an exception occurs are likely
63    /// to repeatedly result in the same exception.
64    #[inline]
65    fn next(&mut self) -> Option<Self::Item> {
66        Borrowed::from(&*self).next()
67    }
68
69    #[cfg(not(Py_LIMITED_API))]
70    fn size_hint(&self) -> (usize, Option<usize>) {
71        let hint = unsafe { ffi::PyObject_LengthHint(self.as_ptr(), 0) };
72        (hint.max(0) as usize, None)
73    }
74}
75
76impl<'py> Borrowed<'_, 'py, PyIterator> {
77    // TODO: this method is on Borrowed so that &'py PyIterator can use this; once that
78    // implementation is deleted this method should be moved to the `Bound<'py, PyIterator> impl
79    fn next(self) -> Option<PyResult<Bound<'py, PyAny>>> {
80        let py = self.py();
81
82        match unsafe { ffi::PyIter_Next(self.as_ptr()).assume_owned_or_opt(py) } {
83            Some(obj) => Some(Ok(obj)),
84            None => PyErr::take(py).map(Err),
85        }
86    }
87}
88
89impl<'py> IntoIterator for &Bound<'py, PyIterator> {
90    type Item = PyResult<Bound<'py, PyAny>>;
91    type IntoIter = Bound<'py, PyIterator>;
92
93    fn into_iter(self) -> Self::IntoIter {
94        self.clone()
95    }
96}
97
98impl PyTypeCheck for PyIterator {
99    const NAME: &'static str = "Iterator";
100
101    fn type_check(object: &Bound<'_, PyAny>) -> bool {
102        unsafe { ffi::PyIter_Check(object.as_ptr()) != 0 }
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::PyIterator;
109    use crate::exceptions::PyTypeError;
110    use crate::types::{PyAnyMethods, PyDict, PyList, PyListMethods};
111    use crate::{ffi, IntoPyObject, Python};
112
113    #[test]
114    fn vec_iter() {
115        Python::with_gil(|py| {
116            let inst = vec![10, 20].into_pyobject(py).unwrap();
117            let mut it = inst.try_iter().unwrap();
118            assert_eq!(
119                10_i32,
120                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
121            );
122            assert_eq!(
123                20_i32,
124                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
125            );
126            assert!(it.next().is_none());
127        });
128    }
129
130    #[test]
131    fn iter_refcnt() {
132        let (obj, count) = Python::with_gil(|py| {
133            let obj = vec![10, 20].into_pyobject(py).unwrap();
134            let count = obj.get_refcnt();
135            (obj.unbind(), count)
136        });
137
138        Python::with_gil(|py| {
139            let inst = obj.bind(py);
140            let mut it = inst.try_iter().unwrap();
141
142            assert_eq!(
143                10_i32,
144                it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
145            );
146        });
147
148        Python::with_gil(move |py| {
149            assert_eq!(count, obj.get_refcnt(py));
150        });
151    }
152
153    #[test]
154    fn iter_item_refcnt() {
155        Python::with_gil(|py| {
156            let count;
157            let obj = py.eval(ffi::c_str!("object()"), None, None).unwrap();
158            let list = {
159                let list = PyList::empty(py);
160                list.append(10).unwrap();
161                list.append(&obj).unwrap();
162                count = obj.get_refcnt();
163                list
164            };
165
166            {
167                let mut it = list.iter();
168
169                assert_eq!(10_i32, it.next().unwrap().extract::<'_, i32>().unwrap());
170                assert!(it.next().unwrap().is(&obj));
171                assert!(it.next().is_none());
172            }
173            assert_eq!(count, obj.get_refcnt());
174        });
175    }
176
177    #[test]
178    fn fibonacci_generator() {
179        let fibonacci_generator = ffi::c_str!(
180            r#"
181def fibonacci(target):
182    a = 1
183    b = 1
184    for _ in range(target):
185        yield a
186        a, b = b, a + b
187"#
188        );
189
190        Python::with_gil(|py| {
191            let context = PyDict::new(py);
192            py.run(fibonacci_generator, None, Some(&context)).unwrap();
193
194            let generator = py
195                .eval(ffi::c_str!("fibonacci(5)"), None, Some(&context))
196                .unwrap();
197            for (actual, expected) in generator.try_iter().unwrap().zip(&[1, 1, 2, 3, 5]) {
198                let actual = actual.unwrap().extract::<usize>().unwrap();
199                assert_eq!(actual, *expected)
200            }
201        });
202    }
203
204    #[test]
205    fn fibonacci_generator_bound() {
206        use crate::types::any::PyAnyMethods;
207        use crate::Bound;
208
209        let fibonacci_generator = ffi::c_str!(
210            r#"
211def fibonacci(target):
212    a = 1
213    b = 1
214    for _ in range(target):
215        yield a
216        a, b = b, a + b
217"#
218        );
219
220        Python::with_gil(|py| {
221            let context = PyDict::new(py);
222            py.run(fibonacci_generator, None, Some(&context)).unwrap();
223
224            let generator: Bound<'_, PyIterator> = py
225                .eval(ffi::c_str!("fibonacci(5)"), None, Some(&context))
226                .unwrap()
227                .downcast_into()
228                .unwrap();
229            let mut items = vec![];
230            for actual in &generator {
231                let actual = actual.unwrap().extract::<usize>().unwrap();
232                items.push(actual);
233            }
234            assert_eq!(items, [1, 1, 2, 3, 5]);
235        });
236    }
237
238    #[test]
239    fn int_not_iterable() {
240        Python::with_gil(|py| {
241            let x = 5i32.into_pyobject(py).unwrap();
242            let err = PyIterator::from_object(&x).unwrap_err();
243
244            assert!(err.is_instance_of::<PyTypeError>(py));
245        });
246    }
247
248    #[test]
249    #[cfg(feature = "macros")]
250    fn python_class_not_iterator() {
251        use crate::PyErr;
252
253        #[crate::pyclass(crate = "crate")]
254        struct Downcaster {
255            failed: Option<PyErr>,
256        }
257
258        #[crate::pymethods(crate = "crate")]
259        impl Downcaster {
260            fn downcast_iterator(&mut self, obj: &crate::Bound<'_, crate::PyAny>) {
261                self.failed = Some(obj.downcast::<PyIterator>().unwrap_err().into());
262            }
263        }
264
265        // Regression test for 2913
266        Python::with_gil(|py| {
267            let downcaster = crate::Py::new(py, Downcaster { failed: None }).unwrap();
268            crate::py_run!(
269                py,
270                downcaster,
271                r#"
272                    from collections.abc import Sequence
273
274                    class MySequence(Sequence):
275                        def __init__(self):
276                            self._data = [1, 2, 3]
277
278                        def __getitem__(self, index):
279                            return self._data[index]
280
281                        def __len__(self):
282                            return len(self._data)
283
284                    downcaster.downcast_iterator(MySequence())
285                "#
286            );
287
288            assert_eq!(
289                downcaster.borrow_mut(py).failed.take().unwrap().to_string(),
290                "TypeError: 'MySequence' object cannot be converted to 'Iterator'"
291            );
292        });
293    }
294
295    #[test]
296    #[cfg(feature = "macros")]
297    fn python_class_iterator() {
298        #[crate::pyfunction(crate = "crate")]
299        fn assert_iterator(obj: &crate::Bound<'_, crate::PyAny>) {
300            assert!(obj.downcast::<PyIterator>().is_ok())
301        }
302
303        // Regression test for 2913
304        Python::with_gil(|py| {
305            let assert_iterator = crate::wrap_pyfunction!(assert_iterator, py).unwrap();
306            crate::py_run!(
307                py,
308                assert_iterator,
309                r#"
310                    class MyIter:
311                        def __next__(self):
312                            raise StopIteration
313
314                    assert_iterator(MyIter())
315                "#
316            );
317        });
318    }
319
320    #[test]
321    #[cfg(not(Py_LIMITED_API))]
322    fn length_hint_becomes_size_hint_lower_bound() {
323        Python::with_gil(|py| {
324            let list = py.eval(ffi::c_str!("[1, 2, 3]"), None, None).unwrap();
325            let iter = list.try_iter().unwrap();
326            let hint = iter.size_hint();
327            assert_eq!(hint, (3, None));
328        });
329    }
330}
pyo3/types/iterator.rs

pyo3/types/
iterator.rs
