diff --git a/Lib/test/test_iterlen.py b/Lib/test/test_iterlen.py index 51ddb7a4a..f05b716e0 100644 --- a/Lib/test/test_iterlen.py +++ b/Lib/test/test_iterlen.py @@ -148,8 +148,6 @@ class TestDictKeys(TestTemporarilyImmutable, unittest.TestCase): self.it = iter(d) self.mutate = d.popitem - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_immutable_during_iteration(self): super().test_immutable_during_iteration() @@ -160,8 +158,6 @@ class TestDictItems(TestTemporarilyImmutable, unittest.TestCase): self.it = iter(d.items()) self.mutate = d.popitem - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_immutable_during_iteration(self): super().test_immutable_during_iteration() @@ -172,8 +168,6 @@ class TestDictValues(TestTemporarilyImmutable, unittest.TestCase): self.it = iter(d.values()) self.mutate = d.popitem - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_immutable_during_iteration(self): super().test_immutable_during_iteration() diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index 283dcc10d..0a0c12573 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -5,6 +5,7 @@ use std::mem::size_of; use super::pystr::PyStrRef; use super::pytype::PyTypeRef; use super::set::PySet; +use super::IterStatus; use crate::dictdatatype::{self, DictKey}; use crate::exceptions::PyBaseExceptionRef; use crate::function::{FuncArgs, KwArgs, OptionalArg}; @@ -675,6 +676,7 @@ macro_rules! dict_iterator { pub dict: PyDictRef, pub size: dictdatatype::DictSize, pub position: AtomicCell, + pub status: AtomicCell, } impl PyValue for $iter_name { @@ -690,30 +692,44 @@ macro_rules! dict_iterator { position: AtomicCell::new(0), size: dict.size(), dict, + status: AtomicCell::new(IterStatus::Active), } } #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { - self.dict.entries.len_from_entry_index(self.position.load()) + if let IterStatus::Exhausted = self.status.load() { + 0 + } else { + self.dict.entries.len_from_entry_index(self.position.load()) + } } } impl PyIter for $iter_name { #[allow(clippy::redundant_closure_call)] fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - if zelf.dict.entries.has_changed_size(&zelf.size) { - return Err( - vm.new_runtime_error("dictionary changed size during iteration".to_owned()) - ); - } - let mut position = zelf.position.load(); - match zelf.dict.entries.next_entry(&mut position) { - Some((key, value)) => { - zelf.position.store(position); - Ok(($result_fn)(vm, key, value)) + match zelf.status.load() { + IterStatus::Exhausted => Err(vm.new_stop_iteration()), + IterStatus::Active => { + if zelf.dict.entries.has_changed_size(&zelf.size) { + zelf.status.store(IterStatus::Exhausted); + return Err(vm.new_runtime_error( + "dictionary changed size during iteration".to_owned(), + )); + } + let mut position = zelf.position.load(); + match zelf.dict.entries.next_entry(&mut position) { + Some((key, value)) => { + zelf.position.store(position); + Ok(($result_fn)(vm, key, value)) + } + None => { + zelf.status.store(IterStatus::Exhausted); + Err(vm.new_stop_iteration()) + } + } } - None => Err(vm.new_stop_iteration()), } } } @@ -724,6 +740,7 @@ macro_rules! dict_iterator { pub dict: PyDictRef, pub size: dictdatatype::DictSize, pub position: AtomicCell, + pub status: AtomicCell, } impl PyValue for $reverse_iter_name { @@ -739,32 +756,43 @@ macro_rules! dict_iterator { position: AtomicCell::new(1), size: dict.size(), dict, + status: AtomicCell::new(IterStatus::Active), } } #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { - self.dict.entries.len_from_entry_index(self.position.load()) + if let IterStatus::Exhausted = self.status.load() { + 0 + } else { + self.dict.entries.len_from_entry_index(self.position.load()) + } } } impl PyIter for $reverse_iter_name { #[allow(clippy::redundant_closure_call)] fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - if zelf.dict.entries.has_changed_size(&zelf.size) { - return Err( - vm.new_runtime_error("dictionary changed size during iteration".to_owned()) - ); - } - let count = zelf.position.fetch_add(1); - match zelf.dict.len().checked_sub(count) { - Some(mut pos) => { - let (key, value) = zelf.dict.entries.next_entry(&mut pos).unwrap(); - Ok(($result_fn)(vm, key, value)) - } - None => { - zelf.position.store(std::isize::MAX as usize); - Err(vm.new_stop_iteration()) + match zelf.status.load() { + IterStatus::Exhausted => Err(vm.new_stop_iteration()), + IterStatus::Active => { + if zelf.dict.entries.has_changed_size(&zelf.size) { + zelf.status.store(IterStatus::Exhausted); + return Err(vm.new_runtime_error( + "dictionary changed size during iteration".to_owned(), + )); + } + let count = zelf.position.fetch_add(1); + match zelf.dict.len().checked_sub(count) { + Some(mut pos) => { + let (key, value) = zelf.dict.entries.next_entry(&mut pos).unwrap(); + Ok(($result_fn)(vm, key, value)) + } + None => { + zelf.status.store(IterStatus::Exhausted); + Err(vm.new_stop_iteration()) + } + } } } }