Add set.add method

windelbouwman · windelbouwman · commit 57075484369e · 2018-10-17T21:08:46.000+02:00
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -592,6 +592,7 @@ Atom: ast::Expression = {
      }
   },
   "{" <e:TestDict?> "}" => ast::Expression::Dict { elements: e.unwrap_or(Vec::new()) },
+  "{" <e:TestDictComp> "}" => e,
   "{" <e:TestSet> "}" => ast::Expression::Set { elements: e },
   "{" <e:TestSetComp> "}" => e,
   "True" => ast::Expression::True,
@@ -623,6 +624,15 @@ TestDict: Vec<(ast::Expression, ast::Expression)> = {
     }
 };
 
+TestDictComp: ast::Expression = {
+    <e1:DictEntry> <c:CompFor> => {
+        ast::Expression::Comprehension {
+            kind: Box::new(ast::ComprehensionKind::Dict { key: e1.0, value: e1.1 }),
+            generators: c,
+        }
+    }
+};
+
 DictEntry: (ast::Expression, ast::Expression) = {
     <e1: Test> ":" <e2: Test> => (e1, e2),
 };
diff --git a/tests/snippets/comprehensions.py b/tests/snippets/comprehensions.py
@@ -10,3 +10,11 @@
     (1, 2), (1, 5), (1, 10),
     (2, 2), (2, 5), (2, 10),
     (3, 2), (3, 5), (3, 10)]
+
+v = {b * 2 for b in x}
+# TODO: how to check set equality?
+# assert v == {2, 6, 4}
+
+# TODO:
+#u = {str(b): b-2 for b in x}
+
diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs
@@ -116,7 +116,7 @@ fn list_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.new_str(s))
 }
 
-pub fn append(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+pub fn list_append(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     trace!("list.append called with: {:?}", args);
     arg_check!(
         vm,
@@ -151,7 +151,7 @@ fn list_len(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.context().new_int(elements.len() as i32))
 }
 
-fn reverse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+fn list_reverse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     trace!("list.reverse called with: {:?}", args);
     arg_check!(vm, args, required = [(list, Some(vm.ctx.list_type()))]);
     let mut list_obj = list.borrow_mut();
@@ -192,7 +192,7 @@ pub fn init(context: &PyContext) {
     list_type.set_attr("__len__", context.new_rustfunc(list_len));
     list_type.set_attr("__new__", context.new_rustfunc(list_new));
     list_type.set_attr("__repr__", context.new_rustfunc(list_repr));
-    list_type.set_attr("append", context.new_rustfunc(append));
+    list_type.set_attr("append", context.new_rustfunc(list_append));
     list_type.set_attr("clear", context.new_rustfunc(clear));
-    list_type.set_attr("reverse", context.new_rustfunc(reverse));
+    list_type.set_attr("reverse", context.new_rustfunc(list_reverse));
 }
diff --git a/vm/src/obj/objset.rs b/vm/src/obj/objset.rs
@@ -30,6 +30,24 @@ pub fn sequence_to_hashmap(iterable: &Vec<PyObjectRef>) -> HashMap<usize, PyObje
     elements
 }
 
+fn set_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    trace!("set.add called with: {:?}", args);
+    arg_check!(
+        vm,
+        args,
+        required = [(s, Some(vm.ctx.set_type())), (item, None)]
+    );
+    let mut mut_obj = s.borrow_mut();
+
+    if let PyObjectKind::Set { ref mut elements } = mut_obj.kind {
+        let key = item.get_id();
+        elements.insert(key, item.clone());
+        Ok(vm.get_none())
+    } else {
+        Err(vm.new_type_error("set.add is called with no list".to_string()))
+    }
+}
+
 /* Create a new object of sub-type of set */
 fn set_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
@@ -115,4 +133,5 @@ pub fn init(context: &PyContext) {
     set_type.set_attr("__len__", context.new_rustfunc(set_len));
     set_type.set_attr("__new__", context.new_rustfunc(set_new));
     set_type.set_attr("__repr__", context.new_rustfunc(set_repr));
+    set_type.set_attr("add", context.new_rustfunc(set_add));
 }
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -371,7 +371,7 @@ impl VirtualMachine {
                         self.ctx.new_int(lineno.get_row() as i32),
                         self.ctx.new_str(run_obj_name.clone()),
                     ]);
-                    objlist::append(
+                    objlist::list_append(
                         self,
                         PyFuncArgs {
                             args: vec![traceback, pos],
@@ -946,7 +946,7 @@ impl VirtualMachine {
                 }
             }
             bytecode::Instruction::MapAdd { i } => {
-                let dict_obj = self.nth_value(*i);
+                let dict_obj = self.nth_value(*i + 1);
                 let key = self.pop_value();
                 let value = self.pop_value();
                 match self.call_method(&dict_obj, "__setitem__", vec![key, value]) {
