Add dict iterator __length_hint__

RustPython · coolreader18 · Jan 10, 2020 · Dec 29, 2019 · Dec 29, 2019 · Dec 29, 2019
commit a26e5f23f86f75c9087c3b2a4ab546a56ff46ff7
diff --git a/tests/snippets/dict.py b/tests/snippets/dict.py
@@ -254,3 +254,16 @@ def __eq__(self, other):
 assert not {}.__ne__({})
 assert {}.__ne__({'a':'b'})
 assert {}.__ne__(1) == NotImplemented
+
+it = iter({0: 1, 2: 3, 4:5, 6:7})
+assert it.__length_hint__() == 4
+next(it)
+assert it.__length_hint__() == 3
+next(it)
+assert it.__length_hint__() == 2
+next(it)
+assert it.__length_hint__() == 1
+next(it)
+assert it.__length_hint__() == 0
+assert_raises(StopIteration, next, it)
+assert it.__length_hint__() == 0
diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs
@@ -204,26 +204,26 @@ impl<T: Clone> Dict<T> {
     }
 
     pub fn next_entry(&self, position: &mut EntryIndex) -> Option<(&PyObjectRef, &T)> {
-        while *position < self.entries.len() {
-            if let Some(DictEntry { key, value, .. }) = &self.entries[*position] {
-                *position += 1;
-                return Some((key, value));
-            }
+        self.entries[*position..].iter().find_map(|entry| {
             *position += 1;
-        }
-        None
+            entry
+                .as_ref()
+                .map(|DictEntry { key, value, .. }| (key, value))
+        })
+    }
+
+    pub fn len_from_entry_index(&self, position: EntryIndex) -> usize {
        self.entries[position..].iter().flatten().count()
     }
 
     pub fn has_changed_size(&self, position: &DictSize) -> bool {
         position.size != self.size || self.entries.len() != position.entries_size
     }
 
-    pub fn keys<'a>(&'a self) -> Box<dyn Iterator<Item = PyObjectRef> + 'a> {
-        Box::new(
-            self.entries
-                .iter()
-                .filter_map(|v| v.as_ref().map(|v| v.key.clone())),
-        )
+    pub fn keys<'a>(&'a self) -> impl Iterator<Item = PyObjectRef> + 'a {
+        self.entries
+            .iter()
+            .filter_map(|v| v.as_ref().map(|v| v.key.clone()))
     }
 
     /// Lookup the index for the given key.

diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs
@@ -545,6 +545,14 @@ macro_rules! dict_iterator {
             fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
                 zelf
             }
+
+            #[pymethod(name = "__length_hint__")]
+            fn length_hint(&self, _vm: &VirtualMachine) -> usize {
+                self.dict
+                    .entries
+                    .borrow()
+                    .len_from_entry_index(self.position.get())
+            }
         }
 
         impl PyValue for $iter_name {