Implement dict __setitem__

Author: windelbouwman

tests/snippets/comprehensions.py

@@ -15,8 +15,9 @@
 # TODO: how to check set equality?
 # assert v == {2, 6, 4}
 
-# TODO:
-#u = {str(b): b-2 for b in x}
+u = {str(b): b-2 for b in x}
+assert u['3'] == 1
+assert u['1'] == -1
 
 y = [a+2 for a in x if a % 2]
 print(y)

vm/src/builtins.rs

@@ -95,10 +95,8 @@ fn builtin_bin(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.new_str(s))
 }
 
-// builtin_bool
 // builtin_breakpoint
 // builtin_bytearray
-// builtin_bytes
 // builtin_callable
 
 fn builtin_chr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -219,7 +217,6 @@ fn builtin_exec(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_filter
-// builtin_float
 // builtin_format
 // builtin_frozenset
 
@@ -275,7 +272,6 @@ fn builtin_id(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_input
-// builtin_int
 
 fn builtin_isinstance(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(obj, None), (typ, None)]);
@@ -325,8 +321,6 @@ fn builtin_len(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
-// builtin_list
-
 fn builtin_locals(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args);
     Ok(vm.get_locals())
@@ -412,7 +406,6 @@ fn builtin_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
-// builtin_object
 // builtin_oct
 // builtin_open
 
@@ -498,14 +491,12 @@ fn builtin_range(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
-// builtin_repr
 fn builtin_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(obj, None)]);
     vm.to_repr(obj)
 }
 // builtin_reversed
 // builtin_round
-// builtin_set
 
 fn builtin_setattr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
@@ -570,12 +561,10 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("list"), ctx.list_type());
     dict.insert(String::from("locals"), ctx.new_rustfunc(builtin_locals));
     dict.insert(String::from("map"), ctx.new_rustfunc(builtin_map));
-
     dict.insert(String::from("max"), ctx.new_rustfunc(builtin_max));
     dict.insert(String::from("min"), ctx.new_rustfunc(builtin_min));
-
+    dict.insert(String::from("object"), ctx.object());
     dict.insert(String::from("ord"), ctx.new_rustfunc(builtin_ord));
-
     dict.insert(String::from("next"), ctx.new_rustfunc(builtin_next));
     dict.insert(String::from("pow"), ctx.new_rustfunc(builtin_pow));
     dict.insert(String::from("print"), ctx.new_rustfunc(builtin_print));
@@ -586,7 +575,6 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("str"), ctx.str_type());
     dict.insert(String::from("tuple"), ctx.tuple_type());
     dict.insert(String::from("type"), ctx.type_type());
-    dict.insert(String::from("object"), ctx.object());
 
     // Exceptions:
     dict.insert(

vm/src/frame.rs

@@ -247,14 +247,14 @@ impl Frame {
                 Ok(None)
             }
             bytecode::Instruction::BuildMap { size, unpack } => {
-                let mut elements = HashMap::new();
+                let mut elements: HashMap<String, PyObjectRef> = HashMap::new();
                 for _x in 0..*size {
                     let obj = self.pop_value();
                     if *unpack {
                         // Take all key-value pairs from the dict:
                         let dict_elements = objdict::get_elements(&obj);
-                        for (key, obj) in dict_elements {
-                            elements.insert(key, obj);
+                        for (key, obj) in dict_elements.iter() {
+                            elements.insert(key.clone(), obj.clone());
                         }
                     } else {
                         // XXX: Currently, we only support String keys, so we have to unwrap the
@@ -471,7 +471,8 @@ impl Frame {
                     bytecode::CallType::Ex(has_kwargs) => {
                         let kwargs = if *has_kwargs {
                             let kw_dict = self.pop_value();
-                            objdict::get_elements(&kw_dict).into_iter().collect()
+                            let dict_elements = objdict::get_elements(&kw_dict).clone();
+                            dict_elements.into_iter().collect()
                         } else {
                             vec![]
                         };

vm/src/obj/objdict.rs

@@ -6,17 +6,9 @@ use super::super::vm::VirtualMachine;
 use super::objstr;
 use super::objtype;
 use num_bigint::ToBigInt;
+use std::cell::{Ref, RefMut};
 use std::collections::HashMap;
-
-pub fn _set_item(
-    vm: &mut VirtualMachine,
-    _d: PyObjectRef,
-    _idx: PyObjectRef,
-    _obj: PyObjectRef,
-) -> PyResult {
-    // TODO: Implement objdict::set_item
-    Ok(vm.get_none())
-}
+use std::ops::{Deref, DerefMut};
 
 pub fn new(dict_type: PyObjectRef) -> PyObjectRef {
     PyObject::new(
@@ -27,12 +19,28 @@ pub fn new(dict_type: PyObjectRef) -> PyObjectRef {
     )
 }
 
-pub fn get_elements(obj: &PyObjectRef) -> HashMap<String, PyObjectRef> {
-    if let PyObjectKind::Dict { elements } = &obj.borrow().kind {
-        elements.clone()
-    } else {
-        panic!("Cannot extract dict elements");
-    }
+pub fn get_elements<'a>(
+    obj: &'a PyObjectRef,
+) -> impl Deref<Target = HashMap<String, PyObjectRef>> + 'a {
+    Ref::map(obj.borrow(), |py_obj| {
+        if let PyObjectKind::Dict { ref elements } = py_obj.kind {
+            elements
+        } else {
+            panic!("Cannot extract dict elements");
+        }
+    })
+}
+
+fn get_mut_elements<'a>(
+    obj: &'a PyObjectRef,
+) -> impl DerefMut<Target = HashMap<String, PyObjectRef>> + 'a {
+    RefMut::map(obj.borrow_mut(), |py_obj| {
+        if let PyObjectKind::Dict { ref mut elements } = py_obj.kind {
+            elements
+        } else {
+            panic!("Cannot extract dict elements");
+        }
+    })
 }
 
 fn dict_new(_vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -50,13 +58,13 @@ fn dict_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
     let elements = get_elements(o);
     let mut str_parts = vec![];
-    for elem in elements {
+    for elem in elements.iter() {
         let s = vm.to_repr(&elem.1)?;
         let value_str = objstr::get_value(&s);
         str_parts.push(format!("{}: {}", elem.0, value_str));
     }
 
-    let s = format!("{{ {} }}", str_parts.join(", "));
+    let s = format!("{{{}}}", str_parts.join(", "));
     Ok(vm.new_str(s))
 }
 
@@ -80,7 +88,7 @@ pub fn dict_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.new_bool(false))
 }
 
-pub fn dict_delitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+fn dict_delitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
@@ -94,18 +102,34 @@ pub fn dict_delitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let needle = objstr::get_value(&needle);
 
     // Delete the item:
-    let mut dict_obj = dict.borrow_mut();
-    if let PyObjectKind::Dict { ref mut elements } = dict_obj.kind {
-        match elements.remove(&needle) {
-            Some(_) => Ok(vm.get_none()),
-            None => Err(vm.new_value_error(format!("Key not found: {}", needle))),
-        }
-    } else {
-        panic!("Cannot extract dict elements");
+    let mut elements = get_mut_elements(dict);
+    match elements.remove(&needle) {
+        Some(_) => Ok(vm.get_none()),
+        None => Err(vm.new_value_error(format!("Key not found: {}", needle))),
     }
 }
 
-pub fn dict_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+fn dict_setitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [
+            (dict, Some(vm.ctx.dict_type())),
+            (needle, Some(vm.ctx.str_type())),
+            (value, None)
+        ]
+    );
+
+    // What we are looking for:
+    let needle = objstr::get_value(&needle);
+
+    // Delete the item:
+    let mut elements = get_mut_elements(dict);
+    elements.insert(needle, value.clone());
+    Ok(vm.get_none())
+}
+
+fn dict_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
@@ -143,4 +167,5 @@ pub fn init(context: &PyContext) {
     dict_type.set_attr("__getitem__", context.new_rustfunc(dict_getitem));
     dict_type.set_attr("__new__", context.new_rustfunc(dict_new));
     dict_type.set_attr("__repr__", context.new_rustfunc(dict_repr));
+    dict_type.set_attr("__setitem__", context.new_rustfunc(dict_setitem));
 }

vm/src/obj/objset.rs

@@ -98,13 +98,17 @@ fn set_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(o, Some(vm.ctx.set_type()))]);
 
     let elements = get_elements(o);
-    let mut str_parts = vec![];
-    for elem in elements.values() {
-        let part = vm.to_repr(elem)?;
-        str_parts.push(objstr::get_value(&part));
-    }
+    let s = if elements.len() == 0 {
+        "set()".to_string()
+    } else {
+        let mut str_parts = vec![];
+        for elem in elements.values() {
+            let part = vm.to_repr(elem)?;
+            str_parts.push(objstr::get_value(&part));
+        }
 
-    let s = format!("{{ {} }}", str_parts.join(", "));
+        format!("{{{}}}", str_parts.join(", "))
+    };
     Ok(vm.new_str(s))
 }
 

vm/src/obj/objtype.rs

@@ -209,7 +209,7 @@ pub fn get_attributes(obj: &PyObjectRef) -> HashMap<String, PyObjectRef> {
         } = &bc.borrow().kind
         {
             let elements = objdict::get_elements(dict);
-            for (name, value) in elements {
+            for (name, value) in elements.iter() {
                 attributes.insert(name.to_string(), value.clone());
             }
         }
@@ -218,7 +218,7 @@ pub fn get_attributes(obj: &PyObjectRef) -> HashMap<String, PyObjectRef> {
     // Get instance attributes:
     if let PyObjectKind::Instance { dict } = &obj.borrow().kind {
         let elements = objdict::get_elements(dict);
-        for (name, value) in elements {
+        for (name, value) in elements.iter() {
             attributes.insert(name.to_string(), value.clone());
         }
     }

vm/src/stdlib/json.rs

@@ -62,7 +62,7 @@ impl<'s> serde::Serialize for PyObjectSerializer<'s> {
         } else if objtype::isinstance(self.pyobject, &self.ctx.dict_type()) {
             let elements = objdict::get_elements(self.pyobject);
             let mut map = serializer.serialize_map(Some(elements.len()))?;
-            for (key, e) in elements {
+            for (key, e) in elements.iter() {
                 map.serialize_entry(&key, &self.clone_with_object(&e))?;
             }
             map.end()