Skip to content

Commit ec5fd55

Browse files
cthulahoopscthulahoops
committed
Implement IntoIterator for PyDictRef.
1 parent 31fd293 commit ec5fd55

6 files changed

Lines changed: 63 additions & 47 deletions

File tree

‎vm/src/dictdatatype.rs‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -131,18 +131,6 @@ impl<T: Clone> Dict<T> {
131131
None
132132
}
133133

134-
pub fn iter_items(&self) -> impl Iterator<Item = (PyObjectRef, T)> + '_ {
135-
self.entries
136-
.iter()
137-
.filter(|e| e.is_some())
138-
.map(|e| e.as_ref().unwrap())
139-
.map(|e| (e.key.clone(), e.value.clone()))
140-
}
141-
142-
pub fn get_items(&self) -> Vec<(PyObjectRef, T)> {
143-
self.iter_items().collect()
144-
}
145-
146134
/// Lookup the index for the given key.
147135
fn lookup(&self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<LookupResult> {
148136
let hash_value = calc_hash(vm, key)?;

‎vm/src/frame.rs‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -393,9 +393,8 @@ impl Frame {
393393
// Take all key-value pairs from the dict:
394394
let dict: PyDictRef =
395395
obj.downcast().expect("Need a dictionary to build a map.");
396-
let dict_elements = dict.get_key_value_pairs();
397-
for (key, value) in dict_elements.iter() {
398-
map_obj.set_item(key.clone(), value.clone(), vm).unwrap();
396+
for (key, value) in dict {
397+
map_obj.set_item(key, value, vm).unwrap();
399398
}
400399
}
401400
} else {
@@ -636,8 +635,7 @@ impl Frame {
636635
let kwargs = if *has_kwargs {
637636
let kw_dict: PyDictRef =
638637
self.pop_value().downcast().expect("Kwargs must be a dict.");
639-
let dict_elements = kw_dict.get_key_value_pairs();
640-
dict_elements
638+
kw_dict
641639
.into_iter()
642640
.map(|elem| (objstr::get_value(&elem.0), elem.1))
643641
.collect()
@@ -862,8 +860,8 @@ impl Frame {
862860

863861
// Grab all the names from the module and put them in the context
864862
if let Some(dict) = &module.dict {
865-
for (k, v) in dict.get_key_value_pairs().iter() {
866-
self.scope.store_name(&vm, &objstr::get_value(k), v.clone());
863+
for (k, v) in dict {
864+
self.scope.store_name(&vm, &objstr::get_value(&k), v);
867865
}
868866
}
869867
Ok(None)
@@ -1230,8 +1228,7 @@ impl fmt::Debug for Frame {
12301228
.collect::<String>();
12311229
let dict = self.scope.get_locals();
12321230
let local_str = dict
1233-
.get_key_value_pairs()
1234-
.iter()
1231+
.into_iter()
12351232
.map(|elem| format!("\n {:?} = {:?}", elem.0, elem.1))
12361233
.collect::<String>();
12371234
write!(

‎vm/src/obj/objdict.rs‎

Lines changed: 54 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@ use std::fmt;
33

44
use crate::function::{KwArgs, OptionalArg};
55
use crate::pyobject::{
6-
IdProtocol, IntoPyObject, ItemProtocol, PyAttributes, PyContext, PyObjectRef, PyRef, PyResult,
7-
PyValue,
6+
IntoPyObject, ItemProtocol, PyAttributes, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
87
};
98
use crate::vm::{ReprGuard, VirtualMachine};
109

@@ -63,9 +62,8 @@ impl PyDictRef {
6362
if let OptionalArg::Present(dict_obj) = dict_obj {
6463
let dicted: PyResult<PyDictRef> = dict_obj.clone().downcast();
6564
if let Ok(dict_obj) = dicted {
66-
let mut dict_borrowed = dict.borrow_mut();
67-
for (key, value) in dict_obj.entries.borrow().iter_items() {
68-
dict_borrowed.insert(vm, &key, value)?;
65+
for (key, value) in dict_obj {
66+
dict.borrow_mut().insert(vm, &key, value)?;
6967
}
7068
} else {
7169
let iter = objiter::get_iter(vm, &dict_obj)?;
@@ -105,7 +103,7 @@ impl PyDictRef {
105103
fn repr(self, vm: &VirtualMachine) -> PyResult<String> {
106104
let s = if let Some(_guard) = ReprGuard::enter(self.as_object()) {
107105
let mut str_parts = vec![];
108-
for (key, value) in self.get_key_value_pairs() {
106+
for (key, value) in self {
109107
let key_repr = vm.to_repr(&key)?;
110108
let value_repr = vm.to_repr(&value)?;
111109
str_parts.push(format!("{}: {}", key_repr.value, value_repr.value));
@@ -146,10 +144,6 @@ impl PyDictRef {
146144
PyDictItems::new(self)
147145
}
148146

149-
pub fn get_key_value_pairs(&self) -> Vec<(PyObjectRef, PyObjectRef)> {
150-
self.entries.borrow().get_items()
151-
}
152-
153147
fn inner_setitem(
154148
self,
155149
key: PyObjectRef,
@@ -194,24 +188,17 @@ impl PyDictRef {
194188

195189
fn update(
196190
self,
197-
mut dict_obj: OptionalArg<PyObjectRef>,
191+
dict_obj: OptionalArg<PyObjectRef>,
198192
kwargs: KwArgs,
199193
vm: &VirtualMachine,
200194
) -> PyResult<()> {
201-
if let OptionalArg::Present(ref other) = dict_obj {
202-
if self.is(other) {
203-
// Updating yourself is a noop, and this avoids a borrow error
204-
dict_obj = OptionalArg::Missing;
205-
}
206-
}
207-
208195
PyDictRef::merge(&self.entries, dict_obj, kwargs, vm)
209196
}
210197

211198
/// Take a python dictionary and convert it to attributes.
212199
pub fn to_attributes(self) -> PyAttributes {
213200
let mut attrs = PyAttributes::new();
214-
for (key, value) in self.get_key_value_pairs() {
201+
for (key, value) in self {
215202
let key = objstr::get_value(&key);
216203
attrs.insert(key, value);
217204
}
@@ -247,6 +234,54 @@ impl ItemProtocol for PyDictRef {
247234
}
248235
}
249236

237+
// Implement IntoIterator so that we can easily iterate dictionaries from rust code.
238+
impl IntoIterator for PyDictRef {
239+
type Item = (PyObjectRef, PyObjectRef);
240+
type IntoIter = DictIterator;
241+
242+
fn into_iter(self) -> Self::IntoIter {
243+
DictIterator::new(self)
244+
}
245+
}
246+
247+
impl IntoIterator for &PyDictRef {
248+
type Item = (PyObjectRef, PyObjectRef);
249+
type IntoIter = DictIterator;
250+
251+
fn into_iter(self) -> Self::IntoIter {
252+
DictIterator::new(self.clone())
253+
}
254+
}
255+
256+
pub struct DictIterator {
257+
dict: PyDictRef,
258+
position: Cell<usize>,
259+
}
260+
261+
impl DictIterator {
262+
pub fn new(dict: PyDictRef) -> DictIterator {
263+
DictIterator {
264+
dict,
265+
position: Cell::new(0),
266+
}
267+
}
268+
}
269+
270+
impl Iterator for DictIterator {
271+
type Item = (PyObjectRef, PyObjectRef);
272+
273+
fn next(&mut self) -> Option<Self::Item> {
274+
self.dict
275+
.entries
276+
.borrow()
277+
.next_entry(self.position.get())
278+
.map(|(new_position, key, value)| {
279+
self.position.set(new_position);
280+
(key.clone(), value.clone())
281+
})
282+
}
283+
}
284+
250285
macro_rules! dict_iterator {
251286
( $name: ident, $iter_name: ident, $class: ident, $iter_class: ident, $class_name: literal, $iter_class_name: literal, $result_fn: expr) => {
252287
#[pyclass(name = $class_name, __inside_vm)]

‎vm/src/obj/objmodule.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,7 @@ impl PyValue for PyModule {
1717
impl PyModuleRef {
1818
fn dir(self: PyModuleRef, vm: &VirtualMachine) -> PyResult {
1919
if let Some(dict) = &self.into_object().dict {
20-
let keys = dict
21-
.get_key_value_pairs()
22-
.iter()
23-
.map(|(k, _v)| k.clone())
24-
.collect();
20+
let keys = dict.into_iter().map(|(k, _v)| k.clone()).collect();
2521
Ok(vm.ctx.new_list(keys))
2622
} else {
2723
panic!("Modules should definitely have a dict.");

‎vm/src/obj/objobject.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ pub fn get_attributes(obj: &PyObjectRef) -> PyAttributes {
237237

238238
// Get instance attributes:
239239
if let Some(dict) = &obj.dict {
240-
for (key, value) in dict.get_key_value_pairs() {
240+
for (key, value) in dict.into_iter() {
241241
attributes.insert(key.to_string(), value.clone());
242242
}
243243
}

‎vm/src/stdlib/json.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ impl<'s> serde::Serialize for PyObjectSerializer<'s> {
6464
serialize_seq_elements(serializer, &elements)
6565
} else if objtype::isinstance(self.pyobject, &self.vm.ctx.dict_type()) {
6666
let dict: PyDictRef = self.pyobject.clone().downcast().unwrap();
67-
let pairs = dict.get_key_value_pairs();
67+
let pairs: Vec<(PyObjectRef, PyObjectRef)> = dict.into_iter().collect();
6868
let mut map = serializer.serialize_map(Some(pairs.len()))?;
6969
for (key, e) in pairs.iter() {
7070
map.serialize_entry(&self.clone_with_object(key), &self.clone_with_object(&e))?;

0 commit comments

Comments
 (0)