Merge pull request #2875 from DimitrisJim/retry_set_ops

youknowone · web-flow · commit 634a6c243529 · 2021-08-14T00:27:20.000+09:00
Retrying contains, remove, discard with a frozenset.
diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py
@@ -379,8 +379,6 @@ class TestSet(TestJointOps, unittest.TestCase):
     thetype = set
     basetype = set
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_contains(self):
         super().test_contains()
 
@@ -440,8 +438,6 @@ def test_add(self):
         self.assertEqual(self.s, dup)
         self.assertRaises(TypeError, self.s.add, [])
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_remove(self):
         self.s.remove('a')
         self.assertNotIn('a', self.s)
@@ -464,8 +460,6 @@ def test_remove_keyerror_unpacking(self):
             else:
                 self.fail()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_remove_keyerror_set(self):
         key = self.thetype([3, 4])
         try:
@@ -477,8 +471,6 @@ def test_remove_keyerror_set(self):
         else:
             self.fail()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_discard(self):
         self.s.discard('a')
         self.assertNotIn('a', self.s)
diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs
@@ -93,7 +93,7 @@ impl PySetInner {
     }
 
     fn contains(&self, needle: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
-        self.content.contains(vm, needle)
+        self.retry_op_with_frozenset(needle, vm, |needle, vm| self.content.contains(vm, needle))
     }
 
     fn compare(
@@ -213,11 +213,11 @@ impl PySetInner {
     }
 
     fn remove(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
-        self.content.delete(vm, item)
+        self.retry_op_with_frozenset(&item, vm, |item, vm| self.content.delete(vm, item.clone()))
     }
 
     fn discard(&self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
-        self.content.delete_if_exists(vm, item)
+        self.retry_op_with_frozenset(item, vm, |item, vm| self.content.delete_if_exists(vm, item))
     }
 
     fn clear(&self) {
@@ -285,6 +285,42 @@ impl PySetInner {
     fn hash(&self, vm: &VirtualMachine) -> PyResult<PyHash> {
         crate::utils::hash_iter_unordered(self.content.keys().iter(), vm)
     }
+
+    // Run operation, on failure, if item is a set/set subclass, convert it
+    // into a frozenset and try the operation again. Propagates original error
+    // on failure to convert and restores item in KeyError on failure (remove).
+    fn retry_op_with_frozenset<T, F>(
+        &self,
+        item: &PyObjectRef,
+        vm: &VirtualMachine,
+        op: F,
+    ) -> PyResult<T>
+    where
+        F: Fn(&PyObjectRef, &VirtualMachine) -> PyResult<T>,
+    {
+        op(item, vm).or_else(|original_err| {
+            item.payload_if_subclass::<PySet>(vm)
+                // Keep original error around.
+                .ok_or(original_err)
+                .and_then(|set| {
+                    op(
+                        &PyFrozenSet {
+                            inner: set.inner.copy(),
+                        }
+                        .into_object(vm),
+                        vm,
+                    )
+                    // If operation raised KeyError, report original set (set.remove)
+                    .map_err(|op_err| {
+                        if op_err.isinstance(&vm.ctx.exceptions.key_error) {
+                            vm.new_key_error(item.clone())
+                        } else {
+                            op_err
+                        }
+                    })
+                })
+        })
+    }
 }
 
 fn extract_set(obj: &PyObjectRef) -> Option<&PySetInner> {
