diff --git a/crates/vm/src/builtins/iter.rs b/crates/vm/src/builtins/iter.rs index e9f1516b5..4e29df583 100644 --- a/crates/vm/src/builtins/iter.rs +++ b/crates/vm/src/builtins/iter.rs @@ -184,17 +184,27 @@ impl PySequenceIterator { } #[pymethod] - fn __length_hint__(&self, vm: &VirtualMachine) -> PyObjectRef { - let internal = self.internal.lock(); - if let IterStatus::Active(obj) = &internal.status { - let seq = obj.sequence_unchecked(); - seq.length(vm).map_or_else( - |_| vm.ctx.not_implemented(), - |x| PyInt::from(x).into_pyobject(vm), - ) - } else { - PyInt::from(0).into_pyobject(vm) - } + fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult { + vm.with_recursion("in __length_hint__", || { + let (obj, position) = { + let internal = self.internal.lock(); + match &internal.status { + IterStatus::Active(obj) => (Some(obj.clone()), internal.position), + IterStatus::Exhausted => (None, 0), + } + }; + if let Some(obj) = obj { + let seq = obj.sequence_unchecked(); + match seq.length_opt(vm) { + Some(len) => { + len.map(|len| PyInt::from(len.saturating_sub(position)).into_pyobject(vm)) + } + None => Ok(vm.ctx.not_implemented()), + } + } else { + Ok(PyInt::from(0).into_pyobject(vm)) + } + }) } #[pymethod] diff --git a/extra_tests/snippets/builtin_iter.py b/extra_tests/snippets/builtin_iter.py new file mode 100644 index 000000000..02d469a47 --- /dev/null +++ b/extra_tests/snippets/builtin_iter.py @@ -0,0 +1,71 @@ +import queue +import threading + + +def make_iterator(): + holder = {} + + class Evil: + def __getitem__(self, index): + if index == 0: + return 0 + raise IndexError + + def __len__(self): + return holder["it"].__length_hint__() + + obj = Evil() + holder["it"] = iter(obj) + return holder["it"] + + +it = make_iterator() +q = queue.Queue() + + +def run(): + try: + it.__length_hint__() + except Exception as exc: # noqa: BLE001 + q.put(exc) + else: + q.put(None) + + +t = threading.Thread(target=run, daemon=True) +t.start() +t.join(1) + +assert not t.is_alive(), "iterator.__length_hint__ deadlocked" +err = q.get_nowait() +assert isinstance(err, RecursionError) + + +class NoLen: + def __getitem__(self, index): + if index < 3: + return index + raise IndexError + + +no_len_it = iter(NoLen()) +assert no_len_it.__length_hint__() is NotImplemented +next(no_len_it) +assert no_len_it.__length_hint__() is NotImplemented + + +class Seq: + def __init__(self): + self.items = [1, 2, 3] + + def __getitem__(self, index): + return self.items[index] + + def __len__(self): + return len(self.items) + + +seq_it = iter(Seq()) +assert seq_it.__length_hint__() == 3 +next(seq_it) +assert seq_it.__length_hint__() == 2