Fully write test in C

Author: eendebakpt

Lib/test/test_capi/test_set.py

@@ -222,6 +222,10 @@ def test_clear(self):
 
 
 class TestPySet_Add(unittest.TestCase):
+    def test_pyset_add(self):
+        # Run C-level tests for PySet_Add
+        _testcapi.test_pyset_add()
+
     def test_set(self):
         # Test the PySet_Add c-api for set objects
         s = set()

Modules/_testcapimodule.c

@@ -2434,6 +2434,153 @@ test_critical_sections(PyObject *module, PyObject *Py_UNUSED(args))
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_pyset_add_exact_set(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    // Test: Adding to a regular set
+    PyObject *set = PySet_New(NULL);
+    if (set == NULL) {
+        return NULL;
+    }
+    PyObject *one = PyLong_FromLong(1);
+    assert(one);
+
+    if (PySet_Add(set, one) < 0) {
+        Py_DECREF(set);
+        return raiseTestError(self, "test_pyset_add",
+                              "PySet_Add to empty set failed");
+    }
+    if (PySet_Size(set) != 1) {
+        Py_DECREF(set);
+        return raiseTestError(self, "test_pyset_add",
+                              "set size should be 1 after adding one element");
+    }
+    if (PySet_Contains(set, one) != 1) {
+        Py_DECREF(set);
+        return raiseTestError(self, "test_pyset_add",
+                              "set should contain the added element");
+    }
+    Py_DECREF(set);
+
+    // Test: Adding unhashable item should raise TypeError
+    set = PySet_New(NULL);
+    if (set == NULL) {
+        return NULL;
+    }
+    PyObject *unhashable = PyList_New(0);
+    if (unhashable == NULL) {
+        Py_DECREF(set);
+        return NULL;
+    }
+    if (PySet_Add(set, unhashable) != -1) {
+        Py_DECREF(unhashable);
+        Py_DECREF(set);
+        return raiseTestError(self, "test_pyset_add",
+                              "PySet_Add with unhashable should fail");
+    }
+    if (!PyErr_ExceptionMatches(PyExc_TypeError)) {
+        Py_DECREF(unhashable);
+        Py_DECREF(set);
+        return raiseTestError(self, "test_pyset_add",
+                              "PySet_Add with unhashable should raise TypeError");
+    }
+    PyErr_Clear();
+    Py_DECREF(unhashable);
+    Py_DECREF(set);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_pyset_add_frozenset(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyObject *one = PyLong_FromLong(1);
+    assert(one);
+
+    // Test: Adding to uniquely-referenced frozenset should succeed
+    PyObject *frozenset = PyFrozenSet_New(NULL);
+    if (frozenset == NULL) {
+        return NULL;
+    }
+
+    // frozenset is uniquely referenced here, so PySet_Add should work
+    if (PySet_Add(frozenset, one) < 0) {
+        Py_DECREF(frozenset);
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "PySet_Add to uniquely-referenced frozenset failed");
+    }
+    Py_DECREF(frozenset);
+
+    // Test: Adding to non-uniquely-referenced frozenset should raise SystemError
+    frozenset = PyFrozenSet_New(NULL);
+    if (frozenset == NULL) {
+        return NULL;
+    }
+    Py_INCREF(frozenset);  // Make it non-uniquely referenced
+
+    if (PySet_Add(frozenset, one) != -1) {
+        Py_DECREF(frozenset);
+        Py_DECREF(frozenset);
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "PySet_Add to non-uniquely-referenced frozenset should fail");
+    }
+    if (!PyErr_ExceptionMatches(PyExc_SystemError)) {
+        Py_DECREF(frozenset);
+        Py_DECREF(frozenset);
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "PySet_Add to non-uniquely-referenced frozenset should raise SystemError");
+    }
+    PyErr_Clear();
+    Py_DECREF(frozenset);
+    Py_DECREF(frozenset);
+
+    // Test: GC tracking - frozenset with only immutable items should not be tracked
+    frozenset = PyFrozenSet_New(NULL);
+    if (frozenset == NULL) {
+        return NULL;
+    }
+    if (PySet_Add(frozenset, one) < 0) {
+        Py_DECREF(frozenset);
+        return NULL;
+    }
+    if (PyObject_GC_IsTracked(frozenset)) {
+        Py_DECREF(frozenset);
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "frozenset with only int should not be GC tracked");
+    }
+    Py_DECREF(frozenset);
+
+    // Test: GC tracking - frozenset with tracked object should be tracked
+    frozenset = PyFrozenSet_New(NULL);
+    if (frozenset == NULL) {
+        return NULL;
+    }
+
+    PyObject *tracked_obj = PyErr_NewException("hashable_and_tracked", NULL, NULL);
+    if (tracked_obj == NULL) {
+        Py_DECREF(frozenset);
+        return NULL;
+    }
+    if (!PyObject_GC_IsTracked(tracked_obj)) {
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "test object should be tracked");
+    }
+    if (PySet_Add(frozenset, tracked_obj) < 0) {
+        Py_DECREF(frozenset);
+        Py_DECREF(tracked_obj);
+        return NULL;
+    }
+
+    if (!PyObject_GC_IsTracked(frozenset)) {
+        Py_DECREF(frozenset);
+        Py_DECREF(tracked_obj);
+        return raiseTestError(self, "test_pyset_add_frozenset",
+                              "frozenset with with GC tracked object should be tracked");
+    }
+
+    Py_RETURN_NONE;
+}
+
 
 // Used by `finalize_thread_hang`.
 #if defined(_POSIX_THREADS) && !defined(__wasi__)
@@ -2685,6 +2832,8 @@ static PyMethodDef TestMethods[] = {
     {"test_weakref_capi", test_weakref_capi, METH_NOARGS},
     {"function_set_warning", function_set_warning, METH_NOARGS},
     {"test_critical_sections", test_critical_sections, METH_NOARGS},
+    {"test_pyset_add_exact_set", test_pyset_add_exact_set, METH_NOARGS},
+    {"test_pyset_add_frozenset", test_pyset_add_frozenset, METH_NOARGS},
     {"finalize_thread_hang", finalize_thread_hang, METH_O, NULL},
     {"test_atexit", test_atexit, METH_NOARGS},
     {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},