Merge pull request #2859 from DimitrisJim/fix_faulty_length_hint

Fix length hint for dictionary iterators.
This commit is contained in:
Jim Fasarakis-Hilliard authored and GitHub committed 2021-08-12 00:20:06 +03:00
commit 5982e8c11c
2 files changed
+55 -33

No files matched your search

-6
View File
@@ -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()
+55 -27
View File
@@ -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<usize>,
pub status: AtomicCell<IterStatus>,
}
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<Self>, 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<usize>,
pub status: AtomicCell<IterStatus>,
}
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<Self>, 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())
}
}
}
}
}