diff --git a/Lib/test/test_csv.py b/Lib/test/test_csv.py index 3e86af0f8c5..1108846f544 100644 --- a/Lib/test/test_csv.py +++ b/Lib/test/test_csv.py @@ -568,7 +568,6 @@ def test_roundtrip_escaped_unquoted_newlines(self): self.assertEqual(row, rows[i]) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: Error not raised def test_reader_reentrant_iterator(self): # gh-145105: re-entering the reader from the iterator must not crash. class ReentrantIter: diff --git a/crates/stdlib/src/csv.rs b/crates/stdlib/src/csv.rs index cb7cecff416..f4ccda55412 100644 --- a/crates/stdlib/src/csv.rs +++ b/crates/stdlib/src/csv.rs @@ -411,6 +411,7 @@ mod _csv { skipinitialspace: options.get_skipinitialspace(), delimiter: options.get_delimiter(), line_num: 0, + generation: 0, }), dialect: options.result(vm)?, }) @@ -961,6 +962,7 @@ mod _csv { skipinitialspace: bool, delimiter: u8, line_num: u64, + generation: u64, } #[pyclass(no_attr, module = "_csv", name = "reader", traverse)] @@ -1057,8 +1059,18 @@ mod _csv { impl IterNext for Reader { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let string = raise_if_stop!(zelf.iter.next(vm)?); - let string = string.downcast::().map_err(|obj| { + let generation = zelf.state.lock().generation; + let string_obj = raise_if_stop!(zelf.iter.next(vm)?); + let mut state = zelf.state.lock(); + if state.generation != generation { + return Err(new_csv_error( + vm, + "iterator has already advanced the reader", + )); + } + state.generation += 1; + + let string = string_obj.downcast::().map_err(|obj| { new_csv_error( vm, format!( @@ -1071,7 +1083,6 @@ mod _csv { if input.is_empty() || input.starts_with(b"\n") { return Ok(PyIterReturn::Return(vm.ctx.new_list(vec![]).into())); } - let mut state = zelf.state.lock(); let ReadState { buffer, output_ends, @@ -1079,6 +1090,7 @@ mod _csv { skipinitialspace, delimiter, line_num, + generation: _, } = &mut *state; let mut input_offset = 0;