Skip to content

Commit 8c4c636

Browse files
authored
fix(itertools): add re-entrancy guard to tee object (RustPython#5931)
* fix(itertools): add re-entrancy guard to tee object * apply feedback PyRwLock -> PyMutex & remove AtomicCell<bool> lock field
1 parent 18d7c1b commit 8c4c636

File tree

2 files changed

+11
-8
lines changed

2 files changed

+11
-8
lines changed

Lib/test/test_itertools.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1761,8 +1761,6 @@ def test_tee_del_backward(self):
17611761
del forward, backward
17621762
raise
17631763

1764-
# TODO: RUSTPYTHON
1765-
@unittest.expectedFailure
17661764
def test_tee_reenter(self):
17671765
class I:
17681766
first = True

vm/src/stdlib/itertools.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1184,23 +1184,28 @@ mod decl {
11841184
#[derive(Debug)]
11851185
struct PyItertoolsTeeData {
11861186
iterable: PyIter,
1187-
values: PyRwLock<Vec<PyObjectRef>>,
1187+
values: PyMutex<Vec<PyObjectRef>>,
11881188
}
11891189

11901190
impl PyItertoolsTeeData {
11911191
fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyResult<PyRc<Self>> {
11921192
Ok(PyRc::new(Self {
11931193
iterable,
1194-
values: PyRwLock::new(vec![]),
1194+
values: PyMutex::new(vec![]),
11951195
}))
11961196
}
11971197

11981198
fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult<PyIterReturn> {
1199-
if self.values.read().len() == index {
1200-
let result = raise_if_stop!(self.iterable.next(vm)?);
1201-
self.values.write().push(result);
1199+
let Some(mut values) = self.values.try_lock() else {
1200+
return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
1201+
};
1202+
1203+
if values.len() == index {
1204+
let obj = raise_if_stop!(self.iterable.next(vm)?);
1205+
values.push(obj);
12021206
}
1203-
Ok(PyIterReturn::Return(self.values.read()[index].clone()))
1207+
1208+
Ok(PyIterReturn::Return(values[index].clone()))
12041209
}
12051210
}
12061211

0 commit comments

Comments
 (0)