Issue #2042: add optimized versions of set.remove() and set.discard()

Cython/Compiler/Builtin.py

@@ -328,7 +328,10 @@ builtin_types_table = [
     ("set",       "PySet_Type",    [BuiltinMethod("__contains__",  "TO",   "b", "PySequence_Contains"),
                                     BuiltinMethod("clear",   "T",  "r", "PySet_Clear"),
                                     # discard() and remove() have a special treatment for unhashable values
-#                                    BuiltinMethod("discard", "TO", "r", "PySet_Discard"),
+                                    BuiltinMethod("discard", "TO", "r", "__Pyx_PySet_Discard",
+                                                  utility_code=UtilityCode.load("py_set_discard", "Optimize.c")),
+                                    BuiltinMethod("remove",  "TO", "r", "__Pyx_PySet_Remove",
+                                                  utility_code=UtilityCode.load("py_set_remove", "Optimize.c")),
                                     # update is actually variadic (see Github issue #1645)
 #                                    BuiltinMethod("update",     "TO", "r", "__Pyx_PySet_Update",
 #                                                  utility_code=UtilityCode.load_cached("PySet_Update", "Builtins.c")),

Cython/Utility/Optimize.c

@@ -397,6 +397,69 @@ static CYTHON_INLINE int __Pyx_dict_iter_next(
 }
 
 
+/////////////// py_set_discard.proto ///////////////
+
+static CYTHON_INLINE int __Pyx_PySet_Discard(PyObject *set, PyObject *key); /*proto*/
+
+/////////////// py_set_discard ///////////////
+
+static CYTHON_INLINE int __Pyx_PySet_Discard(PyObject *set, PyObject *key) {
+    PyObject *tmpkey;
+    int rv;
+
+    rv = PySet_Discard(set, key);
+    /* Convert *key* to frozenset if necessary */
+    if (rv < 0) {
+        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError))
+            return -1;
+        PyErr_Clear();
+        tmpkey = PyFrozenSet_New(key);
+        if (tmpkey == NULL)
+            return -1;
+        rv = PySet_Discard(set, tmpkey);
+        Py_DECREF(tmpkey);
+    }
+    return rv;
+}
+
+
+/////////////// py_set_remove.proto ///////////////
+
+static CYTHON_INLINE int __Pyx_PySet_Remove(PyObject *set, PyObject *key); /*proto*/
+
+/////////////// py_set_remove ///////////////
+
+static CYTHON_INLINE int __Pyx_PySet_Remove(PyObject *set, PyObject *key) {
+    PyObject *tmpkey;
+    int rv;
+
+    rv = PySet_Discard(set, key);
+
+    /* Convert *key* to frozenset if necessary */
+    if (rv < 0) {
+        if (!PySet_Check(key) || !PyErr_ExceptionMatches(PyExc_TypeError))
+            return -1;
+        PyErr_Clear();
+        tmpkey = PyFrozenSet_New(key);
+        if (tmpkey == NULL)
+            return -1;
+        rv = PySet_Discard(set, tmpkey);
+        Py_DECREF(tmpkey);
+    }
+    if (rv == 0) {
+        /* Not found */
+        PyObject *tup;
+        tup = PyTuple_Pack(1, key);
+        if (!tup)
+            return -1;
+        PyErr_SetObject(PyExc_KeyError, tup);
+        Py_DECREF(tup);
+        return -1;
+    }
+    return 0;
+}
+
+
 /////////////// unicode_iter.proto ///////////////
 
 static CYTHON_INLINE int __Pyx_init_unicode_iteration(

tests/run/set.pyx

@@ -65,6 +65,32 @@ def test_set_add():
     return s1
 
 
+def test_set_contains(v):
+    """
+    >>> test_set_contains(1)
+    True
+    >>> test_set_contains(2)
+    False
+    >>> test_set_contains(frozenset([1, 2, 3]))
+    True
+    >>> test_set_contains(frozenset([1, 2]))
+    False
+    >>> test_set_contains(set([1, 2, 3]))
+    True
+    >>> test_set_contains(set([1, 2]))
+    False
+    >>> try: test_set_contains([1, 2])
+    ... except TypeError: pass
+    ... else: print("NOT RAISED!")
+    """
+    cdef set s1
+    s1 = set()
+    s1.add(1)
+    s1.add('a')
+    s1.add(frozenset([1, 2, 3]))
+    return v in s1
+
+
 def test_set_update(v=None):
     """
     >>> type(test_set_update()) is set