gh-155102: Preserve set discard/remove comparison errors

Author: BHUVANSH855

Lib/test/test_set.py

@@ -471,6 +471,29 @@ def test_discard(self):
         self.assertNotIn(self.thetype(self.word), s)
         s.discard(self.thetype(self.word))
 
+    def test_discard_remove_hashable_set_subclass_eq_typeerror(self):
+        class Bad:
+            def __hash__(self):
+                return 1
+
+            def __eq__(self, other):
+                raise TypeError("boom from __eq__")
+
+        class HashableSet(set):
+            def __hash__(self):
+                return 1
+
+        probe = HashableSet()
+
+        with self.assertRaisesRegex(TypeError, "boom from __eq__"):
+            probe in {Bad()}
+
+        with self.assertRaisesRegex(TypeError, "boom from __eq__"):
+            {Bad()}.discard(probe)
+
+        with self.assertRaisesRegex(TypeError, "boom from __eq__"):
+            {Bad()}.remove(probe)
+
     def test_pop(self):
         for i in range(len(self.s)):
             elem = self.s.pop()

Objects/setobject.c

@@ -2653,20 +2653,22 @@ set_remove_impl(PySetObject *so, PyObject *key)
 {
     int rv;
 
-    rv = set_discard_key(so, key);
-    if (rv < 0) {
-        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError))
+    Py_hash_t hash = PyObject_Hash(key);
+    if (hash == -1) {
+        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError)) {
+            set_unhashable_type(key);
             return NULL;
+        }
         PyErr_Clear();
-        Py_hash_t hash;
         Py_BEGIN_CRITICAL_SECTION(key);
         hash = frozenset_hash_impl(key);
         Py_END_CRITICAL_SECTION();
-        rv = set_discard_entry(so, key, hash);
-        if (rv < 0)
-            return NULL;
     }
 
+    rv = set_discard_entry(so, key, hash);
+    if (rv < 0)
+        return NULL;
+
     if (rv == DISCARD_NOTFOUND) {
         _PyErr_SetKeyError(key);
         return NULL;
@@ -2693,19 +2695,21 @@ set_discard_impl(PySetObject *so, PyObject *key)
 {
     int rv;
 
-    rv = set_discard_key(so, key);
-    if (rv < 0) {
-        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError))
+    Py_hash_t hash = PyObject_Hash(key);
+    if (hash == -1) {
+        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError)) {
+            set_unhashable_type(key);
             return NULL;
+        }
         PyErr_Clear();
-        Py_hash_t hash;
         Py_BEGIN_CRITICAL_SECTION(key);
         hash = frozenset_hash_impl(key);
         Py_END_CRITICAL_SECTION();
-        rv = set_discard_entry(so, key, hash);
-        if (rv < 0)
-            return NULL;
     }
+
+    rv = set_discard_entry(so, key, hash);
+    if (rv < 0)
+        return NULL;
     Py_RETURN_NONE;
 }