Add PyWeakref_GetRef() function

vstinner · vstinner · commit ad8bae847cda · 2023-06-21T13:27:31.000+02:00
diff --git a/docs/api.rst b/docs/api.rst
@@ -31,6 +31,9 @@ Python 3.13
 
    See `PyImport_AddModuleRef() documentation <https://docs.python.org/dev/c-api/import.html#c.PyImport_AddModuleRef>`__.
 
+.. c:function:: int PyWeakref_GetRef(PyObject *ref, PyObject **pobj)
+
+   See `PyWeakref_GetRef() documentation <https://docs.python.org/dev/c-api/weakref.html#c.PyWeakref_GetRef>`__.
 
 Python 3.12
 -----------
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -1,6 +1,7 @@
 Changelog
 =========
 
+* 2023-06-21: Add ``PyWeakref_GetRef()`` function.
 * 2023-06-20: Add ``PyImport_AddModuleRef()`` function.
 * 2022-11-15: Add experimental operations to the ``upgrade_pythoncapi.py``
   script: ``Py_NewRef``, ``Py_CLEAR`` and ``Py_SETREF``.
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -570,6 +570,7 @@ PyCode_GetCellvars(PyCodeObject *code)
 #  endif
 #endif
 
+
 // gh-105922 added PyImport_AddModuleRef() to Python 3.13.0a1
 #if PY_VERSION_HEX < 0x030D00A0
 PYCAPI_COMPAT_STATIC_INLINE(PyObject*)
@@ -580,6 +581,33 @@ PyImport_AddModuleRef(const char *name)
 #endif
 
 
+// gh-105927 added PyWeakref_GetRef() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D0000
+PYCAPI_COMPAT_STATIC_INLINE(int)
+PyWeakref_GetRef(PyObject *ref, PyObject **pobj)
+{
+    PyObject *obj;
+    if (ref != NULL && !PyWeakref_Check(ref)) {
+        *pobj = NULL;
+        PyErr_SetString(PyExc_TypeError, "expected a weakref");
+        return -1;
+    }
+    obj = PyWeakref_GetObject(ref);
+    if (obj == NULL) {
+        // SystemError if ref is NULL
+        *pobj = NULL;
+        return -1;
+    }
+    if (obj == Py_None) {
+        *pobj = NULL;
+        return 0;
+    }
+    *pobj = Py_NewRef(obj);
+    return 0;
+}
+#endif
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -668,6 +668,78 @@ test_import(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+static void
+gc_collect(void)
+{
+#if defined(PYPY_VERSION)
+    PyObject *mod = PyImport_ImportModule("gc");
+    assert(mod != NULL);
+
+    PyObject *res = PyObject_CallMethod(mod, "collect", NULL);
+    Py_DECREF(mod);
+    assert(res != NULL);
+    Py_DECREF(res);
+#else
+    PyGC_Collect();
+#endif
+}
+
+
+static PyObject *
+test_weakref(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    // Create a new heap type, create an instance of this type, and delete the
+    // type. This object supports weak references.
+    PyObject *new_type = PyObject_CallFunction((PyObject*)&PyType_Type,
+                                               "s(){}", "TypeName");
+    if (new_type == NULL) {
+        return NULL;
+    }
+    PyObject *obj = PyObject_CallNoArgs(new_type);
+    Py_DECREF(new_type);
+    if (obj == NULL) {
+        return NULL;
+    }
+    Py_ssize_t refcnt = Py_REFCNT(obj);
+
+    // create a weak reference
+    PyObject *weakref = PyWeakref_NewRef(obj, NULL);
+    if (weakref == NULL) {
+        return NULL;
+    }
+
+    // test PyWeakref_GetRef(), reference is alive
+    PyObject *ref1;
+    assert(PyWeakref_GetRef(weakref, &ref1) == 0);
+    assert(ref1 == obj);
+    assert(Py_REFCNT(obj) == (refcnt + 1));
+    Py_DECREF(ref1);
+
+    // delete the referenced object
+    assert(Py_REFCNT(obj) == 1);
+    Py_DECREF(obj);
+    gc_collect();
+
+    // test PyWeakref_GetRef(), reference is dead
+    PyObject *ref2 = Py_True;  // marker to check that value was set
+    assert(PyWeakref_GetRef(weakref, &ref2) == 0);
+    assert(ref2 == NULL);
+
+    // test PyWeakref_GetRef(), invalid type
+    PyObject *invalid_weakref = Py_None;
+    assert(!PyErr_Occurred());
+    PyObject *ref3 = Py_True;  // marker to check that value was set
+    assert(PyWeakref_GetRef(invalid_weakref, &ref3) == -1);
+    assert(PyErr_ExceptionMatches(PyExc_TypeError));
+    assert(ref3 == NULL);
+    PyErr_Clear();
+
+    Py_DECREF(weakref);
+
+    Py_RETURN_NONE;
+}
+
+
 static struct PyMethodDef methods[] = {
     {"test_object", test_object, METH_NOARGS, _Py_NULL},
     {"test_py_is", test_py_is, METH_NOARGS, _Py_NULL},
@@ -687,6 +759,7 @@ static struct PyMethodDef methods[] = {
 #endif
     {"test_api_casts", test_api_casts, METH_NOARGS, _Py_NULL},
     {"test_import", test_import, METH_NOARGS, _Py_NULL},
+    {"test_weakref", test_weakref, METH_NOARGS, _Py_NULL},
     {_Py_NULL, _Py_NULL, 0, _Py_NULL}
 };
 
