Add set.symmetric_difference_update

palaviv · palaviv · commit 21b6616e8e3a · 2019-02-21T19:33:20.000+02:00
diff --git a/tests/snippets/set.py b/tests/snippets/set.py
@@ -108,3 +108,8 @@ def __hash__(self):
 a.difference_update([3,4,5])
 assert a == set([1,2])
 assert_raises(TypeError, lambda: a.difference_update(1))
+
+a = set([1,2,3])
+a.symmetric_difference_update([3,4,5])
+assert a == set([1,2,4,5])
+assert_raises(TypeError, lambda: a.difference_update(1))
diff --git a/vm/src/obj/objset.rs b/vm/src/obj/objset.rs
@@ -492,6 +492,35 @@ fn set_combine_update_inner(
     }
 }
 
+fn set_symmetric_difference_update(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.set_type())), (iterable, None)]
+    );
+
+    let mut mut_obj = zelf.borrow_mut();
+
+    match mut_obj.payload {
+        PyObjectPayload::Set { ref mut elements } => {
+            let elements_original = elements.clone();
+            let iterator = objiter::get_iter(vm, iterable)?;
+            while let Ok(v) = vm.call_method(&iterator, "__next__", vec![]) {
+                insert_into_set(vm, elements, &v)?;
+            }
+            for element in elements_original.iter() {
+                let value = vm.call_method(iterable, "__contains__", vec![element.1.clone()])?;
+                if objbool::get_value(&value) {
+                    elements.remove(&element.0.clone());
+                }
+            }
+
+            Ok(vm.get_none())
+        }
+        _ => Err(vm.new_type_error("".to_string())),
+    }
+}
+
 fn frozenset_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(o, Some(vm.ctx.frozenset_type()))]);
 
@@ -574,6 +603,11 @@ pub fn init(context: &PyContext) {
         "difference_update",
         context.new_rustfunc(set_difference_update),
     );
+    context.set_attr(
+        &set_type,
+        "symmetric_difference_update",
+        context.new_rustfunc(set_symmetric_difference_update),
+    );
 
     let frozenset_type = &context.frozenset_type;
 
