gh-110525: Add tests for internal set CAPI

Lib/test/test_capi/test_set.py

@@ -2,8 +2,9 @@
 
 from test.support import import_helper
 
-# Skip this test if the _testcapi module isn't available.
+# Skip this test if the _testcapi or _testinternalcapi modules aren't available.
 _testcapi = import_helper.import_module('_testcapi')
+_testinternalcapi = import_helper.import_module('_testinternalcapi')
 
 class set_subclass(set):
     pass
@@ -12,13 +13,15 @@ class frozenset_subclass(frozenset):
     pass
 
 
-class TestSetCAPI(unittest.TestCase):
+class BaseSetTests:
     def assertImmutable(self, action, *args):
         self.assertRaises(SystemError, action, frozenset(), *args)
         self.assertRaises(SystemError, action, frozenset({1}), *args)
         self.assertRaises(SystemError, action, frozenset_subclass(), *args)
         self.assertRaises(SystemError, action, frozenset_subclass({1}), *args)
 
+
+class TestSetCAPI(BaseSetTests, unittest.TestCase):
     def test_set_check(self):
         check = _testcapi.set_check
         self.assertTrue(check(set()))
@@ -213,3 +216,39 @@ def test_clear(self):
             clear(object())
         self.assertImmutable(clear)
         # CRASHES: clear(NULL)
+
+
+class TestInternalCAPI(BaseSetTests, unittest.TestCase):
+    def test_set_update(self):
+        update = _testinternalcapi.set_update
+        for cls in (set, set_subclass):
+            for it in ('ab', ('a', 'b'), ['a', 'b'],
+                       set('ab'), set_subclass('ab'),
+                       frozenset('ab'), frozenset_subclass('ab')):
+                with self.subTest(cls=cls, it=it):
+                    instance = cls()
+                    self.assertEqual(update(instance, it), 0)
+                    self.assertEqual(instance, {'a', 'b'})
+                    instance = cls(it)
+                    self.assertEqual(update(instance, it), 0)
+                    self.assertEqual(instance, {'a', 'b'})
+            with self.assertRaisesRegex(TypeError, 'object is not iterable'):
+                update(cls(), 1)
+            with self.assertRaisesRegex(TypeError, "unhashable type: 'dict'"):
+                update(cls(), [{}])
+        with self.assertRaises(SystemError):
+            update(object(), 'ab')
+        self.assertImmutable(update, 'ab')
+        # CRASHES: update(NULL, object())
+        # CRASHES: update(instance, NULL)
+        # CRASHES: update(NULL, NULL)
+
+    def test_set_next_entry(self):
+        set_next = _testinternalcapi.set_next_entry
+        for cls in (set, set_subclass, frozenset, frozenset_subclass):
+            with self.subTest(cls=cls):
+                res = set_next(cls([1, 2, 3]))
+                self.assertEqual(set(res), {(1, 1, 1), (2, 2, 2), (3, 3, 3)})
+        with self.assertRaises(SystemError):
+            set_next(object())
+        # CRASHES: set_next(NULL)

Modules/Setup.stdlib.in

@@ -158,7 +158,7 @@
 @MODULE_XXSUBTYPE_TRUE@xxsubtype xxsubtype.c
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
-@MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c
+@MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c _testinternalcapi/set.c
 @MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/set.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/pyos.c _testcapi/immortal.c _testcapi/heaptype_relative.c _testcapi/gc.c
 @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c
 @MODULE__TESTCLINIC_LIMITED_TRUE@_testclinic_limited _testclinic_limited.c

Modules/_testinternalcapi.c

@@ -1602,6 +1602,9 @@ module_exec(PyObject *module)
     if (_PyTestInternalCapi_Init_PyTime(module) < 0) {
         return 1;
     }
+    if (_PyTestInternalCapi_Init_Set(module) < 0) {
+        return 1;
+    }
 
     if (PyModule_Add(module, "SIZEOF_PYGC_HEAD",
                         PyLong_FromSsize_t(sizeof(PyGC_Head))) < 0) {

Modules/_testinternalcapi/parts.h

@@ -12,5 +12,6 @@
 
 int _PyTestInternalCapi_Init_Lock(PyObject *module);
 int _PyTestInternalCapi_Init_PyTime(PyObject *module);
+int _PyTestInternalCapi_Init_Set(PyObject *module);
 
 #endif // Py_TESTINTERNALCAPI_PARTS_H

Modules/_testinternalcapi/set.c

@@ -0,0 +1,83 @@
+#include "parts.h"
+#include "../_testcapi/util.h"  // NULLABLE, RETURN_INT
+
+#include "pycore_setobject.h"
+
+
+static PyObject *
+set_update(PyObject *self, PyObject *args)
+{
+    PyObject *set, *iterable;
+    if (!PyArg_ParseTuple(args, "OO", &set, &iterable)) {
+        return NULL;
+    }
+    NULLABLE(set);
+    NULLABLE(iterable);
+    RETURN_INT(_PySet_Update(set, iterable));
+}
+
+static PyObject *
+set_next_entry(PyObject *self, PyObject *obj)
+{
+    NULLABLE(obj);
+    PyObject *result = PyList_New(0);
+    if (result == NULL) {
+        return NULL;
+    }
+
+    int set_next = 0;
+    Py_ssize_t i = 0, count = 0;
+    Py_hash_t h;
+    PyObject *item;
+    while ((set_next = _PySet_NextEntry(obj, &i, &item, &h))) {
+        if (set_next == -1 && PyErr_Occurred()) {  // obj is not a set
+            goto error;
+        }
+
+        count++;
+        PyObject *index = PyLong_FromSsize_t(count);
+        if (index == NULL) {
+            goto error;
+        }
+        PyObject *hash = PyLong_FromSize_t((size_t)h);
+        if (hash == NULL) {
+            Py_DECREF(index);
+            goto error;
+        }
+        PyObject *tup = PyTuple_Pack(3, index, item, hash);
+        Py_DECREF(index);
+        Py_DECREF(item);
+        Py_DECREF(hash);
+        if (tup == NULL) {
+            goto error;
+        }
+        int res = PyList_Append(result, tup);
+        Py_DECREF(tup);
+        if (res < 0) {
+            goto error;
+        }
+    }
+    assert(count == PyList_GET_SIZE(result));
+    return result;
+
+error:
+    Py_DECREF(result);
+    return NULL;
+}
+
+
+static PyMethodDef TestMethods[] = {
+    {"set_update", set_update, METH_VARARGS},
+    {"set_next_entry", set_next_entry, METH_O},
+
+    {NULL},
+};
+
+int
+_PyTestInternalCapi_Init_Set(PyObject *m)
+{
+    if (PyModule_AddFunctions(m, TestMethods) < 0) {
+        return -1;
+    }
+    return 0;
+}

PCbuild/_testinternalcapi.vcxproj

@@ -96,6 +96,7 @@
     <ClCompile Include="..\Modules\_testinternalcapi.c" />
     <ClCompile Include="..\Modules\_testinternalcapi\pytime.c" />
     <ClCompile Include="..\Modules\_testinternalcapi\test_lock.c" />
+    <ClCompile Include="..\Modules\_testinternalcapi\set.c" />
   </ItemGroup>
   <ItemGroup>
     <ResourceCompile Include="..\PC\python_nt.rc" />