Skip to content

Commit 1306b71

Browse files
youknowoneCopilotgithub-actions[bot]
authored
__length_hint__ (RustPython#6636)
* Initial plan * fix: prevent iterator length_hint deadlock Co-authored-by: youknowone <69878+youknowone@users.noreply.github.com> * fix * Auto-format: cargo fmt --all --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 6131363 commit 1306b71

2 files changed

Lines changed: 92 additions & 11 deletions

File tree

crates/vm/src/builtins/iter.rs

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -184,17 +184,27 @@ impl PySequenceIterator {
184184
}
185185

186186
#[pymethod]
187-
fn __length_hint__(&self, vm: &VirtualMachine) -> PyObjectRef {
188-
let internal = self.internal.lock();
189-
if let IterStatus::Active(obj) = &internal.status {
190-
let seq = obj.sequence_unchecked();
191-
seq.length(vm).map_or_else(
192-
|_| vm.ctx.not_implemented(),
193-
|x| PyInt::from(x).into_pyobject(vm),
194-
)
195-
} else {
196-
PyInt::from(0).into_pyobject(vm)
197-
}
187+
fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
188+
vm.with_recursion("in __length_hint__", || {
189+
let (obj, position) = {
190+
let internal = self.internal.lock();
191+
match &internal.status {
192+
IterStatus::Active(obj) => (Some(obj.clone()), internal.position),
193+
IterStatus::Exhausted => (None, 0),
194+
}
195+
};
196+
if let Some(obj) = obj {
197+
let seq = obj.sequence_unchecked();
198+
match seq.length_opt(vm) {
199+
Some(len) => {
200+
len.map(|len| PyInt::from(len.saturating_sub(position)).into_pyobject(vm))
201+
}
202+
None => Ok(vm.ctx.not_implemented()),
203+
}
204+
} else {
205+
Ok(PyInt::from(0).into_pyobject(vm))
206+
}
207+
})
198208
}
199209

200210
#[pymethod]
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import queue
2+
import threading
3+
4+
5+
def make_iterator():
6+
holder = {}
7+
8+
class Evil:
9+
def __getitem__(self, index):
10+
if index == 0:
11+
return 0
12+
raise IndexError
13+
14+
def __len__(self):
15+
return holder["it"].__length_hint__()
16+
17+
obj = Evil()
18+
holder["it"] = iter(obj)
19+
return holder["it"]
20+
21+
22+
it = make_iterator()
23+
q = queue.Queue()
24+
25+
26+
def run():
27+
try:
28+
it.__length_hint__()
29+
except Exception as exc: # noqa: BLE001
30+
q.put(exc)
31+
else:
32+
q.put(None)
33+
34+
35+
t = threading.Thread(target=run, daemon=True)
36+
t.start()
37+
t.join(1)
38+
39+
assert not t.is_alive(), "iterator.__length_hint__ deadlocked"
40+
err = q.get_nowait()
41+
assert isinstance(err, RecursionError)
42+
43+
44+
class NoLen:
45+
def __getitem__(self, index):
46+
if index < 3:
47+
return index
48+
raise IndexError
49+
50+
51+
no_len_it = iter(NoLen())
52+
assert no_len_it.__length_hint__() is NotImplemented
53+
next(no_len_it)
54+
assert no_len_it.__length_hint__() is NotImplemented
55+
56+
57+
class Seq:
58+
def __init__(self):
59+
self.items = [1, 2, 3]
60+
61+
def __getitem__(self, index):
62+
return self.items[index]
63+
64+
def __len__(self):
65+
return len(self.items)
66+
67+
68+
seq_it = iter(Seq())
69+
assert seq_it.__length_hint__() == 3
70+
next(seq_it)
71+
assert seq_it.__length_hint__() == 2

0 commit comments

Comments
 (0)