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

pyo3/types/
iterator.rs
