Skip to content

Commit 8b9be05

Browse files
Handle errors on PyIterator when calling size_hint (#5604)
* Handle errors on `PyIterator` when calling `size_hint` * Test `PyNotImplemented` return value as well * Capture unraisable error in test * disable test on `3.7`
1 parent 4f579f1 commit 8b9be05

File tree

2 files changed

+51
-1
lines changed

2 files changed

+51
-1
lines changed

newsfragments/5604.fixed.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Handle errors on `PyIterator` when calling `size_hint`

src/types/iterator.rs

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,15 @@ impl<'py> Iterator for Bound<'py, PyIterator> {
108108

109109
#[cfg(not(Py_LIMITED_API))]
110110
fn size_hint(&self) -> (usize, Option<usize>) {
111+
// SAFETY: `self` is a valid iterator object
111112
let hint = unsafe { ffi::PyObject_LengthHint(self.as_ptr(), 0) };
112-
(hint.max(0) as usize, None)
113+
if hint < 0 {
114+
let py = self.py();
115+
PyErr::fetch(py).write_unraisable(py, Some(self));
116+
(0, None)
117+
} else {
118+
(hint as usize, None)
119+
}
113120
}
114121
}
115122

@@ -144,6 +151,8 @@ mod tests {
144151
#[cfg(all(not(PyPy), Py_3_10))]
145152
use crate::types::PyNone;
146153
use crate::types::{PyAnyMethods, PyDict, PyList, PyListMethods};
154+
#[cfg(all(feature = "macros", Py_3_8, not(Py_LIMITED_API)))]
155+
use crate::PyErr;
147156
use crate::{IntoPyObject, PyTypeInfo, Python};
148157

149158
#[test]
@@ -392,6 +401,46 @@ def fibonacci(target):
392401
});
393402
}
394403

404+
#[test]
405+
#[cfg(all(feature = "macros", Py_3_8, not(Py_LIMITED_API)))]
406+
fn length_hint_error() {
407+
#[crate::pyfunction(crate = "crate")]
408+
fn test_size_hint(obj: &crate::Bound<'_, crate::PyAny>, should_error: bool) {
409+
let iter = obj.cast::<PyIterator>().unwrap();
410+
crate::test_utils::UnraisableCapture::enter(obj.py(), |capture| {
411+
assert_eq!((0, None), iter.size_hint());
412+
assert_eq!(should_error, capture.take_capture().is_some());
413+
});
414+
assert!(PyErr::take(obj.py()).is_none());
415+
}
416+
417+
Python::attach(|py| {
418+
let test_size_hint = crate::wrap_pyfunction!(test_size_hint, py).unwrap();
419+
crate::py_run!(
420+
py,
421+
test_size_hint,
422+
r#"
423+
class NoHintIter:
424+
def __next__(self):
425+
raise StopIteration
426+
427+
def __length_hint__(self):
428+
return NotImplemented
429+
430+
class ErrorHintIter:
431+
def __next__(self):
432+
raise StopIteration
433+
434+
def __length_hint__(self):
435+
raise ValueError("bad hint impl")
436+
437+
test_size_hint(NoHintIter(), False)
438+
test_size_hint(ErrorHintIter(), True)
439+
"#
440+
);
441+
});
442+
}
443+
395444
#[test]
396445
fn test_type_object() {
397446
Python::attach(|py| {

0 commit comments

Comments
 (0)