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

pyo3/types/
iterator.rs
