gh-110525: Cover PySet_Add corner case with frozenset objects (GH-110544)

1 files changed, 37 insertions, 0 deletions

diff --git a/Modules/_testcapi/set.c b/Modules/_testcapi/set.c
index 35e686e..2fbd0ae 100644
--- a/Modules/_testcapi/set.c
+++ b/Modules/_testcapi/set.c
@@ -127,6 +127,41 @@ set_clear(PyObject *self, PyObject *obj)
     RETURN_INT(PySet_Clear(obj));
 }
 
+static PyObject *
+test_frozenset_add_in_capi(PyObject *self, PyObject *Py_UNUSED(obj))
+{
+    // Test that `frozenset` can be used with `PySet_Add`,
+    // when frozenset is just created in CAPI.
+    PyObject *fs = PyFrozenSet_New(NULL);
+    if (fs == NULL) {
+        return NULL;
+    }
+    PyObject *num = PyLong_FromLong(1);
+    if (num == NULL) {
+        goto error;
+    }
+    if (PySet_Add(fs, num) < 0) {
+        goto error;
+    }
+    int contains = PySet_Contains(fs, num);
+    if (contains < 0) {
+        goto error;
+    }
+    else if (contains == 0) {
+        goto unexpected;
+    }
+    Py_DECREF(fs);
+    Py_DECREF(num);
+    Py_RETURN_NONE;
+
+unexpected:
+    PyErr_SetString(PyExc_ValueError, "set does not contain expected value");
+error:
+    Py_DECREF(fs);
+    Py_XDECREF(num);
+    return NULL;
+}
+
 static PyMethodDef test_methods[] = {
     {"set_check", set_check, METH_O},
     {"set_checkexact", set_checkexact, METH_O},
@@ -146,6 +181,8 @@ static PyMethodDef test_methods[] = {
     {"set_pop", set_pop, METH_O},
     {"set_clear", set_clear, METH_O},
 
+    {"test_frozenset_add_in_capi", test_frozenset_add_in_capi, METH_NOARGS},
+
     {NULL},
 };