diff --git a/newsfragments/5604.fixed.md b/newsfragments/5604.fixed.md new file mode 100644 index 00000000000..6bc52a41a55 --- /dev/null +++ b/newsfragments/5604.fixed.md @@ -0,0 +1 @@ +Handle errors on `PyIterator` when calling `size_hint` diff --git a/src/types/iterator.rs b/src/types/iterator.rs index 62ab3bac742..2ab1bda1722 100644 --- a/src/types/iterator.rs +++ b/src/types/iterator.rs @@ -108,8 +108,15 @@ impl<'py> Iterator for Bound<'py, PyIterator> { #[cfg(not(Py_LIMITED_API))] fn size_hint(&self) -> (usize, Option) { + // SAFETY: `self` is a valid iterator object let hint = unsafe { ffi::PyObject_LengthHint(self.as_ptr(), 0) }; - (hint.max(0) as usize, None) + if hint < 0 { + let py = self.py(); + PyErr::fetch(py).write_unraisable(py, Some(self)); + (0, None) + } else { + (hint as usize, None) + } } } @@ -144,6 +151,8 @@ mod tests { #[cfg(all(not(PyPy), Py_3_10))] use crate::types::PyNone; use crate::types::{PyAnyMethods, PyDict, PyList, PyListMethods}; + #[cfg(all(feature = "macros", Py_3_8, not(Py_LIMITED_API)))] + use crate::PyErr; use crate::{IntoPyObject, PyTypeInfo, Python}; #[test] @@ -392,6 +401,46 @@ def fibonacci(target): }); } + #[test] + #[cfg(all(feature = "macros", Py_3_8, not(Py_LIMITED_API)))] + fn length_hint_error() { + #[crate::pyfunction(crate = "crate")] + fn test_size_hint(obj: &crate::Bound<'_, crate::PyAny>, should_error: bool) { + let iter = obj.cast::().unwrap(); + crate::test_utils::UnraisableCapture::enter(obj.py(), |capture| { + assert_eq!((0, None), iter.size_hint()); + assert_eq!(should_error, capture.take_capture().is_some()); + }); + assert!(PyErr::take(obj.py()).is_none()); + } + + Python::attach(|py| { + let test_size_hint = crate::wrap_pyfunction!(test_size_hint, py).unwrap(); + crate::py_run!( + py, + test_size_hint, + r#" + class NoHintIter: + def __next__(self): + raise StopIteration + + def __length_hint__(self): + return NotImplemented + + class ErrorHintIter: + def __next__(self): + raise StopIteration + + def __length_hint__(self): + raise ValueError("bad hint impl") + + test_size_hint(NoHintIter(), False) + test_size_hint(ErrorHintIter(), True) + "# + ); + }); + } + #[test] fn test_type_object() { Python::attach(|py| {