Add _Py_StealRef() and _Py_XStealRef()

vstinner · vstinner · commit 02ea717cdffe · 2021-02-16T09:13:54.000+01:00
diff --git a/README.rst b/README.rst
@@ -152,6 +152,12 @@ Borrow variant
 To ease migration of C extensions to the new C API, a variant is provided
 to return borrowed references rather than strong references::
 
+    // Similar to "Py_INCREF(ob); return ob;"
+    PyObject* _Py_StealRef(PyObject *ob);
+
+    // Similar to "Py_XINCREF(ob); return ob;"
+    PyObject* _Py_XStealRef(PyObject *ob);
+
     // PyThreadState_GetFrame()
     PyFrameObject* _PyThreadState_GetFrameBorrow(PyThreadState *tstate)
 
@@ -290,8 +296,9 @@ Links
 Changelog
 =========
 
-* 2020-11-30: Creation of the upgrade_pythoncapi.py script.
-* 2020-06-04: Creation of the pythoncapi_compat.h header file.
+* 2021-02-16: Add ``_Py_StealRef()`` and ``_Py_XStealRef()`` functions.
+* 2020-11-30: Creation of the ``upgrade_pythoncapi.py`` script.
+* 2020-06-04: Creation of the ``pythoncapi_compat.h`` header file.
 
 
 Examples of projects using pythoncapi_compat.h
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -63,6 +63,28 @@ static inline PyObject* _Py_XNewRef(PyObject *obj)
 #endif
 
 
+// See https://bugs.python.org/issue42522
+#if !defined(_Py_StealRef)
+static inline PyObject* __Py_StealRef(PyObject *obj)
+{
+    Py_DECREF(obj);
+    return obj;
+}
+#define _Py_StealRef(obj) __Py_StealRef(_PyObject_CAST(obj))
+#endif
+
+
+// See https://bugs.python.org/issue42522
+#if !defined(_Py_XStealRef)
+static inline PyObject* __Py_XStealRef(PyObject *obj)
+{
+    Py_XDECREF(obj);
+    return obj;
+}
+#define _Py_XStealRef(obj) __Py_XStealRef(_PyObject_CAST(obj))
+#endif
+
+
 // bpo-39573 added Py_SET_REFCNT() to Python 3.9.0a4
 #if PY_VERSION_HEX < 0x030900A4 && !defined(Py_SET_REFCNT)
 static inline void _Py_SET_REFCNT(PyObject *ob, Py_ssize_t refcnt)
@@ -109,9 +131,7 @@ PyFrame_GetCode(PyFrameObject *frame)
 static inline PyCodeObject*
 _PyFrame_GetCodeBorrow(PyFrameObject *frame)
 {
-    PyCodeObject *code = PyFrame_GetCode(frame);
-    Py_DECREF(code);
-    return code;  // borrowed reference
+    return (PyCodeObject *)_Py_StealRef(PyFrame_GetCode(frame));
 }
 
 
@@ -128,9 +148,7 @@ PyFrame_GetBack(PyFrameObject *frame)
 static inline PyFrameObject*
 _PyFrame_GetBackBorrow(PyFrameObject *frame)
 {
-    PyFrameObject *back = PyFrame_GetBack(frame);
-    Py_XDECREF(back);
-    return back;  // borrowed reference
+    return (PyFrameObject *)_Py_XStealRef(PyFrame_GetBack(frame));
 }
 
 
@@ -158,9 +176,7 @@ PyThreadState_GetFrame(PyThreadState *tstate)
 static inline PyFrameObject*
 _PyThreadState_GetFrameBorrow(PyThreadState *tstate)
 {
-    PyFrameObject *frame = PyThreadState_GetFrame(tstate);
-    Py_XDECREF(frame);
-    return frame;  // borrowed reference
+    return (PyFrameObject *)_Py_XStealRef(PyThreadState_GetFrame(tstate));
 }
 
 
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -23,7 +23,7 @@ test_object(PyObject *self, PyObject *ignored)
     }
     Py_ssize_t refcnt = Py_REFCNT(obj);
 
-    // Py_NewRef() and Py_XNewRef()
+    // Py_NewRef()
     PyObject *ref = Py_NewRef(obj);
     assert(ref == obj);
     assert(Py_REFCNT(obj) == (refcnt + 1));
@@ -53,6 +53,35 @@ test_object(PyObject *self, PyObject *ignored)
 }
 
 
+static PyObject *
+test_steal_ref(PyObject *self, PyObject *ignored)
+{
+    PyObject *obj = PyList_New(0);
+    if (obj == NULL) {
+        return NULL;
+    }
+    Py_ssize_t refcnt = Py_REFCNT(obj);
+
+    // _Py_StealRef()
+    Py_INCREF(obj);
+    PyObject *ref = _Py_StealRef(obj);
+    assert(ref == obj);
+    assert(Py_REFCNT(obj) == refcnt);
+
+    // _Py_XStealRef()
+    Py_INCREF(obj);
+    PyObject *xref = _Py_XStealRef(obj);
+    assert(xref == obj);
+    assert(Py_REFCNT(obj) == refcnt);
+
+    assert(_Py_XStealRef(NULL) == NULL);
+
+    assert(Py_REFCNT(obj) == refcnt);
+    Py_DECREF(obj);
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *
 test_frame(PyObject *self, PyObject *ignored)
 {
@@ -285,6 +314,7 @@ test_module(PyObject *self, PyObject *ignored)
 
 static struct PyMethodDef methods[] = {
     {"test_object", test_object, METH_NOARGS},
+    {"test_steal_ref", test_steal_ref, METH_NOARGS},
     {"test_frame", test_frame, METH_NOARGS},
     {"test_thread_state", test_thread_state, METH_NOARGS},
     {"test_interpreter", test_interpreter, METH_NOARGS},
