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 ffiResult, …>::PyObject_GetIter(arg1: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: Python<'_> = self.py();
71 let mut result: *mut PyObject = std::ptr::null_mut();
72 match unsafe { ffi::PyIter_Send(self.as_ptr(), arg: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: Python<'_> = self.py();
110
111 match unsafe { ffi::PyIter_Next(self.as_ptr()).assume_owned_or_opt(py) } {
112 Some(obj: Bound<'_, PyAny>) => 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(obj: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}
400