fix(itertools): add re-entrancy guard to tee object

ever0de · ever0de · commit aaf18058f825 · 2025-07-10T15:32:35.000+09:00
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
@@ -1761,8 +1761,6 @@ def test_tee_del_backward(self):
             del forward, backward
             raise
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_tee_reenter(self):
         class I:
             first = True
diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
@@ -1185,21 +1185,31 @@ mod decl {
     struct PyItertoolsTeeData {
         iterable: PyIter,
         values: PyRwLock<Vec<PyObjectRef>>,
+        locked: AtomicCell<bool>,
     }
 
     impl PyItertoolsTeeData {
         fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyResult<PyRc<Self>> {
             Ok(PyRc::new(Self {
                 iterable,
                 values: PyRwLock::new(vec![]),
+                locked: AtomicCell::new(false),
             }))
         }
 
         fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult<PyIterReturn> {
             if self.values.read().len() == index {
-                let result = raise_if_stop!(self.iterable.next(vm)?);
-                self.values.write().push(result);
+                if self.locked.swap(true) {
+                    return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
+                }
+
+                let result = self.iterable.next(vm);
+                self.locked.store(false);
+
+                let obj = raise_if_stop!(result?);
+                self.values.write().push(obj);
             }
+
             Ok(PyIterReturn::Return(self.values.read()[index].clone()))
         }
     }
