Merge pull request #3259 from Jack-R-lantern/__contains__

youknowone · web-flow · commit eaae1343563a · 2021-10-10T14:01:11.000+09:00
Add contains in dict_items
diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py
@@ -1300,8 +1300,6 @@ def __eq__(self, other):
         except RuntimeError:  # implementation defined
             pass
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_dictitems_contains_use_after_free(self):
         class X:
             def __eq__(self, other):
diff --git a/extra_tests/snippets/builtin_dict.py b/extra_tests/snippets/builtin_dict.py
@@ -26,3 +26,14 @@
         next(dict_reversed)
 assert 'dict' in dict().__doc__
 
+d = {'a': 123, 'b': 456}
+assert 1 not in d.items()
+assert 'a' not in d.items()
+assert 'a', 123 not in d.items()
+assert () not in d.items()
+assert (1) not in d.items()
+assert ('a') not in d.items()
+assert ('a', 123) in d.items()
+assert ('b', 456) in d.items()
+assert ('a', 123, 3) not in d.items()
+assert ('a', 123, 'b', 456) not in d.items()
diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs
@@ -3,6 +3,7 @@ use super::{
     PyTypeRef,
 };
 use crate::{
+    builtins::PyTuple,
     common::ascii,
     dictdatatype::{self, DictKey},
     function::{ArgIterable, FuncArgs, IntoPyObject, KwArgs, OptionalArg},
@@ -982,6 +983,26 @@ impl PyDictItems {
         let inner = zelf.difference(other, vm)?;
         Ok(PySet { inner })
     }
+
+    #[pymethod(magic)]
+    fn contains(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
+        let needle = match_class! {
+            match needle {
+                tuple @ PyTuple => tuple,
+                _ => return Ok(false),
+            }
+        };
+        if needle.len() != 2 {
+            return Ok(false);
+        }
+        let key = needle.fast_getitem(0);
+        if !zelf.dict().contains(key.clone(), vm)? {
+            return Ok(false);
+        }
+        let value = needle.fast_getitem(1);
+        let found = PyDict::getitem(zelf.dict().clone(), key, vm)?;
+        vm.identical_or_equal(&found, &value)
+    }
 }
 
 pub(crate) fn init(context: &PyContext) {
