merged from other brach

TheAnyKey · TheAnyKey · commit d6192cd927de · 2020-05-07T19:48:23.000Z
diff --git a/tests/snippets/dict_union.py b/tests/snippets/dict_union.py
@@ -0,0 +1,83 @@
+
+import testutils
+
+def test_dunion_ior0():
+    a={1:2,2:3}
+    b={3:4,5:6}
+    a|=b
+
+    assert a == {1:2,2:3,3:4,5:6}, f"wrong value assigned {a=}"
+    assert b == {3:4,5:6}, f"right hand side modified, {b=}"
+
+def test_dunion_or0():
+    a={1:2,2:3}
+    b={3:4,5:6}
+    c=a|b
+
+    assert a == {1:2,2:3}, f"left hand side of non-assignment operator modified {a=}"
+    assert b == {3:4,5:6}, f"right hand side of non-assignment operator modified, {b=}"
+    assert c == {1:2,2:3, 3:4, 5:6}, f"unexpected result of dict union {c=}"
+
+
+def test_dunion_or1():
+    a={1:2,2:3}
+    b={3:4,5:6}
+    c=a.__or__(b)
+
+    assert a == {1:2,2:3}, f"left hand side of non-assignment operator modified {a=}"
+    assert b == {3:4,5:6}, f"right hand side of non-assignment operator modified, {b=}"
+    assert c == {1:2,2:3, 3:4, 5:6}, f"unexpected result of dict union {c=}"
+
+
+def test_dunion_ror0():
+    a={1:2,2:3}
+    b={3:4,5:6}
+    c=b.__ror__(a)
+
+    assert a == {1:2,2:3}, f"left hand side of non-assignment operator modified {a=}"
+    assert b == {3:4,5:6}, f"right hand side of non-assignment operator modified, {b=}"
+    assert c == {1:2,2:3, 3:4, 5:6}, f"unexpected result of dict union {c=}"
+
+
+def test_dunion_other_types():
+    def perf_test_or(other_obj):
+        d={1:2}
+        try:
+            d.__or__(other_obj)
+        except:
+            return True
+        return False
+
+    def perf_test_ior(other_obj):
+        d={1:2}
+        try:
+            d.__ior__(other_obj)
+        except:
+            return True
+        return False
+
+    def perf_test_ror(other_obj):
+        d={1:2}
+        try:
+            d.__ror__(other_obj)
+        except:
+            return True
+        return False
+
+    test_fct={'__or__':perf_test_or, '__ror__':perf_test_ror, '__ior__':perf_test_ior}
+    others=['FooBar', 42, [36], set([19]), ['aa'], None]
+    for tfn,tf in test_fct.items():
+        for other in others:
+            assert tf(other), f"Failed: dict {tfn}, accepted {other}"
+
+
+
+
+testutils.skip_if_unsupported(3,9,test_dunion_ior0)
+testutils.skip_if_unsupported(3,9,test_dunion_or0)
+testutils.skip_if_unsupported(3,9,test_dunion_or1)
+testutils.skip_if_unsupported(3,9,test_dunion_ror0)
+testutils.skip_if_unsupported(3,9,test_dunion_other_types)
+
+
+
diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs
@@ -104,6 +104,17 @@ impl PyDictRef {
         Ok(())
     }
 
+    fn merge_dict(
+        dict: &DictContentType,
+        dict_other: PyDictRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
+        for (key, value) in dict_other {
+            dict.insert(vm, &key, value)?;
+        }
+        Ok(())
+    }
+
     #[pyclassmethod]
     fn fromkeys(
         class: PyClassRef,
@@ -320,6 +331,38 @@ impl PyDictRef {
         PyDictRef::merge(&self.entries, dict_obj, kwargs, vm)
     }
 
+    #[pymethod(name = "__ior__")]
+    fn ior(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        let dicted: Result<PyDictRef, _> = other.clone().downcast();
+        if let Ok(other) = dicted {
+            PyDictRef::merge_dict(&self.entries, other, vm)?;
+            return Ok(self.into_object());
+        }
+        Err(vm.new_type_error("__ior__ not implemented for non-dict type".to_owned()))
+    }
+
+    #[pymethod(name = "__ror__")]
+    fn ror(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyDict> {
+        let dicted: Result<PyDictRef, _> = other.clone().downcast();
+        if let Ok(other) = dicted {
+            let other_cp = other.copy();
+            PyDictRef::merge_dict(&other_cp.entries, self, vm)?;
+            return Ok(other_cp);
+        }
+        Err(vm.new_type_error("__ror__ not implemented for non-dict type".to_owned()))
+    }
+
+    #[pymethod(name = "__or__")]
+    fn or(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyDict> {
+        let dicted: Result<PyDictRef, _> = other.clone().downcast();
+        if let Ok(other) = dicted {
+            let self_cp = self.copy();
+            PyDictRef::merge_dict(&self_cp.entries, other, vm)?;
+            return Ok(self_cp);
+        }
+        Err(vm.new_type_error("__or__ not implemented for non-dict type".to_owned()))
+    }
+
     #[pymethod]
     fn pop(
         self,
