1use crate::{ffi, AsPyPointer, Py, PyAny, PyErr, PyNativeType, PyResult, Python};
2use crate::{PyDowncastError, PyTryFrom};
3
4/// A Python iterator object.
5///
6/// # Examples
7///
8/// ```rust
9/// use pyo3::prelude::*;
10///
11/// # fn main() -> PyResult<()> {
12/// Python::with_gil(|py| -> PyResult<()> {
13/// let list = py.eval("iter([1, 2, 3, 4])", None, None)?;
14/// let numbers: PyResult<Vec<usize>> = list
15/// .iter()?
16/// .map(|i| i.and_then(PyAny::extract::<usize>))
17/// .collect();
18/// let sum: usize = numbers?.iter().sum();
19/// assert_eq!(sum, 10);
20/// Ok(())
21/// })
22/// # }
23/// ```
24#[repr(transparent)]
25pub struct PyIterator(PyAny);
26pyobject_native_type_named!(PyIterator);
27pyobject_native_type_extract!(PyIterator);
28
29impl PyIterator {
30 /// Constructs a `PyIterator` from a Python iterable object.
31 ///
32 /// Equivalent to Python's built-in `iter` function.
33 pub fn from_object(obj: &PyAny) -> PyResult<&PyIterator> {
34 unsafe {
35 obj.py()
36 .from_owned_ptr_or_err(ptr:ffi::PyObject_GetIter(arg1:obj.as_ptr()))
37 }
38 }
39}
40
41impl<'p> Iterator for &'p PyIterator {
42 type Item = PyResult<&'p PyAny>;
43
44 /// Retrieves the next item from an iterator.
45 ///
46 /// Returns `None` when the iterator is exhausted.
47 /// If an exception occurs, returns `Some(Err(..))`.
48 /// Further `next()` calls after an exception occurs are likely
49 /// to repeatedly result in the same exception.
50 fn next(&mut self) -> Option<Self::Item> {
51 let py: Python<'_> = self.0.py();
52
53 match unsafe { py.from_owned_ptr_or_opt(ptr:ffi::PyIter_Next(self.0.as_ptr())) } {
54 Some(obj: &PyAny) => Some(Ok(obj)),
55 None => PyErr::take(py).map(Err),
56 }
57 }
58
59 #[cfg(not(Py_LIMITED_API))]
60 fn size_hint(&self) -> (usize, Option<usize>) {
61 let hint: isize = unsafe { ffi::PyObject_LengthHint(self.0.as_ptr(), arg1:0) };
62 (hint.max(0) as usize, None)
63 }
64}
65
66// PyIter_Check does not exist in the limited API until 3.8
67impl<'v> PyTryFrom<'v> for PyIterator {
68 fn try_from<V: Into<&'v PyAny>>(value: V) -> Result<&'v PyIterator, PyDowncastError<'v>> {
69 let value = value.into();
70 unsafe {
71 if ffi::PyIter_Check(obj:value.as_ptr()) != 0 {
72 Ok(value.downcast_unchecked())
73 } else {
74 Err(PyDowncastError::new(from:value, to:"Iterator"))
75 }
76 }
77 }
78
79 fn try_from_exact<V: Into<&'v PyAny>>(value: V) -> Result<&'v PyIterator, PyDowncastError<'v>> {
80 value.into().downcast()
81 }
82
83 #[inline]
84 unsafe fn try_from_unchecked<V: Into<&'v PyAny>>(value: V) -> &'v PyIterator {
85 let ptr: *const PyIterator = value.into() as *const _ as *const PyIterator;
86 &*ptr
87 }
88}
89
90impl Py<PyIterator> {
91 /// Borrows a GIL-bound reference to the PyIterator. By binding to the GIL lifetime, this
92 /// allows the GIL-bound reference to not require `Python` for any of its methods.
93 pub fn as_ref<'py>(&'py self, _py: Python<'py>) -> &'py PyIterator {
94 let any: *const PyAny = self.as_ptr() as *const PyAny;
95 unsafe { PyNativeType::unchecked_downcast(&*any) }
96 }
97
98 /// Similar to [`as_ref`](#method.as_ref), and also consumes this `Py` and registers the
99 /// Python object reference in PyO3's object storage. The reference count for the Python
100 /// object will not be decreased until the GIL lifetime ends.
101 pub fn into_ref(self, py: Python<'_>) -> &PyIterator {
102 unsafe { py.from_owned_ptr(self.into_ptr()) }
103 }
104}
105
106#[cfg(test)]
107mod tests {
108 use super::PyIterator;
109 use crate::exceptions::PyTypeError;
110 use crate::gil::GILPool;
111 use crate::types::{PyDict, PyList};
112 use crate::{Py, PyAny, Python, ToPyObject};
113
114 #[test]
115 fn vec_iter() {
116 Python::with_gil(|py| {
117 let obj = vec![10, 20].to_object(py);
118 let inst = obj.as_ref(py);
119 let mut it = inst.iter().unwrap();
120 assert_eq!(
121 10_i32,
122 it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
123 );
124 assert_eq!(
125 20_i32,
126 it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
127 );
128 assert!(it.next().is_none());
129 });
130 }
131
132 #[test]
133 fn iter_refcnt() {
134 let (obj, count) = Python::with_gil(|py| {
135 let obj = vec![10, 20].to_object(py);
136 let count = obj.get_refcnt(py);
137 (obj, count)
138 });
139
140 Python::with_gil(|py| {
141 let inst = obj.as_ref(py);
142 let mut it = inst.iter().unwrap();
143
144 assert_eq!(
145 10_i32,
146 it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
147 );
148 });
149
150 Python::with_gil(|py| {
151 assert_eq!(count, obj.get_refcnt(py));
152 });
153 }
154
155 #[test]
156 fn iter_item_refcnt() {
157 Python::with_gil(|py| {
158 let count;
159 let obj = py.eval("object()", None, None).unwrap();
160 let list = {
161 let _pool = unsafe { GILPool::new() };
162 let list = PyList::empty(py);
163 list.append(10).unwrap();
164 list.append(obj).unwrap();
165 count = obj.get_refcnt();
166 list.to_object(py)
167 };
168
169 {
170 let _pool = unsafe { GILPool::new() };
171 let inst = list.as_ref(py);
172 let mut it = inst.iter().unwrap();
173
174 assert_eq!(
175 10_i32,
176 it.next().unwrap().unwrap().extract::<'_, i32>().unwrap()
177 );
178 assert!(it.next().unwrap().unwrap().is(obj));
179 assert!(it.next().is_none());
180 }
181 assert_eq!(count, obj.get_refcnt());
182 });
183 }
184
185 #[test]
186 fn fibonacci_generator() {
187 let fibonacci_generator = r#"
188def fibonacci(target):
189 a = 1
190 b = 1
191 for _ in range(target):
192 yield a
193 a, b = b, a + b
194"#;
195
196 Python::with_gil(|py| {
197 let context = PyDict::new(py);
198 py.run(fibonacci_generator, None, Some(context)).unwrap();
199
200 let generator = py.eval("fibonacci(5)", None, Some(context)).unwrap();
201 for (actual, expected) in generator.iter().unwrap().zip(&[1, 1, 2, 3, 5]) {
202 let actual = actual.unwrap().extract::<usize>().unwrap();
203 assert_eq!(actual, *expected)
204 }
205 });
206 }
207
208 #[test]
209 fn int_not_iterable() {
210 Python::with_gil(|py| {
211 let x = 5.to_object(py);
212 let err = PyIterator::from_object(x.as_ref(py)).unwrap_err();
213
214 assert!(err.is_instance_of::<PyTypeError>(py));
215 });
216 }
217
218 #[test]
219
220 fn iterator_try_from() {
221 Python::with_gil(|py| {
222 let obj: Py<PyAny> = vec![10, 20].to_object(py).as_ref(py).iter().unwrap().into();
223 let iter: &PyIterator = obj.downcast(py).unwrap();
224 assert!(obj.is(iter));
225 });
226 }
227
228 #[test]
229 fn test_as_ref() {
230 Python::with_gil(|py| {
231 let iter: Py<PyIterator> = PyAny::iter(PyList::empty(py)).unwrap().into();
232 let mut iter_ref: &PyIterator = iter.as_ref(py);
233 assert!(iter_ref.next().is_none());
234 })
235 }
236
237 #[test]
238 fn test_into_ref() {
239 Python::with_gil(|py| {
240 let bare_iter = PyAny::iter(PyList::empty(py)).unwrap();
241 assert_eq!(bare_iter.get_refcnt(), 1);
242 let iter: Py<PyIterator> = bare_iter.into();
243 assert_eq!(bare_iter.get_refcnt(), 2);
244 let mut iter_ref = iter.into_ref(py);
245 assert!(iter_ref.next().is_none());
246 assert_eq!(iter_ref.get_refcnt(), 2);
247 })
248 }
249
250 #[test]
251 #[cfg(feature = "macros")]
252 fn python_class_not_iterator() {
253 use crate::PyErr;
254
255 #[crate::pyclass(crate = "crate")]
256 struct Downcaster {
257 failed: Option<PyErr>,
258 }
259
260 #[crate::pymethods(crate = "crate")]
261 impl Downcaster {
262 fn downcast_iterator(&mut self, obj: &PyAny) {
263 self.failed = Some(obj.downcast::<PyIterator>().unwrap_err().into());
264 }
265 }
266
267 // Regression test for 2913
268 Python::with_gil(|py| {
269 let downcaster = Py::new(py, Downcaster { failed: None }).unwrap();
270 crate::py_run!(
271 py,
272 downcaster,
273 r#"
274 from collections.abc import Sequence
275
276 class MySequence(Sequence):
277 def __init__(self):
278 self._data = [1, 2, 3]
279
280 def __getitem__(self, index):
281 return self._data[index]
282
283 def __len__(self):
284 return len(self._data)
285
286 downcaster.downcast_iterator(MySequence())
287 "#
288 );
289
290 assert_eq!(
291 downcaster.borrow_mut(py).failed.take().unwrap().to_string(),
292 "TypeError: 'MySequence' object cannot be converted to 'Iterator'"
293 );
294 });
295 }
296
297 #[test]
298 #[cfg(feature = "macros")]
299 fn python_class_iterator() {
300 #[crate::pyfunction(crate = "crate")]
301 fn assert_iterator(obj: &PyAny) {
302 assert!(obj.downcast::<PyIterator>().is_ok())
303 }
304
305 // Regression test for 2913
306 Python::with_gil(|py| {
307 let assert_iterator = crate::wrap_pyfunction!(assert_iterator, py).unwrap();
308 crate::py_run!(
309 py,
310 assert_iterator,
311 r#"
312 class MyIter:
313 def __next__(self):
314 raise StopIteration
315
316 assert_iterator(MyIter())
317 "#
318 );
319 });
320 }
321
322 #[test]
323 #[cfg(not(Py_LIMITED_API))]
324 fn length_hint_becomes_size_hint_lower_bound() {
325 Python::with_gil(|py| {
326 let list = py.eval("[1, 2, 3]", None, None).unwrap();
327 let iter = list.iter().unwrap();
328 let hint = iter.size_hint();
329 assert_eq!(hint, (3, None));
330 });
331 }
332}
333