Add Borrow variant functions

vstinner · vstinner · commit 12724ddbfe83 · 2020-12-02T00:13:01.000+01:00
diff --git a/README.rst b/README.rst
@@ -35,15 +35,35 @@ Copy the header file in your project and include it using::
 Functions
 =========
 
+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::
+
+    // PyThreadState_GetFrame()
+    PyFrameObject* _PyThreadState_GetFrameBorrow(PyThreadState *tstate)
+
+    // PyFrame_GetCode()
+    PyCodeObject* _PyFrame_GetCodeBorrow(PyFrameObject *frame)
+
+    // PyFrame_GetBack()
+    PyFrameObject* _PyFrame_GetBackBorrow(PyFrameObject *frame)
+
+For example, ``tstate->frame`` can be replaced with
+``_PyThreadState_GetFrameBorrow(tstate)`` to avoid accessing directly
+``PyThreadState.frame`` member.
+
+These functions are only available in ``pythoncapi_compat.h`` and are not
+part of the Python C API.
+
 Python 3.10
 -----------
 
 ::
 
     PyObject* Py_NewRef(PyObject *obj);
     PyObject* Py_XNewRef(PyObject *obj);
-    PyObject* _Py_Borrow(PyObject *obj);
-    PyObject* _Py_XBorrow(PyObject *obj);
 
 Python 3.9
 ----------
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -41,22 +41,6 @@ static inline PyObject* Py_XNewRef(PyObject *obj)
 #endif
 
 
-// bpo-42522
-static inline PyObject* __Py_Borrow(PyObject *obj)
-{
-    Py_DECREF(obj);
-    return obj;
-}
-#define _Py_Borrow(obj) __Py_Borrow(_PyObject_CAST(obj))
-
-static inline PyObject* __Py_XBorrow(PyObject *obj)
-{
-    Py_XDECREF(obj);
-    return obj;
-}
-#define _Py_XBorrow(obj) __Py_XBorrow(_PyObject_CAST(obj))
-
-
 // bpo-39573: Py_TYPE(), Py_REFCNT() and Py_SIZE() can no longer be used
 // as l-value in Python 3.10.
 #if PY_VERSION_HEX < 0x030900A4
@@ -96,6 +80,14 @@ PyFrame_GetCode(PyFrameObject *frame)
 }
 #endif
 
+static inline PyCodeObject*
+_PyFrame_GetCodeBorrow(PyFrameObject *frame)
+{
+    PyCodeObject *code = PyFrame_GetCode(frame);
+    Py_DECREF(code);
+    return code;  // borrowed reference
+}
+
 
 #if PY_VERSION_HEX < 0x030900B1
 static inline PyFrameObject*
@@ -108,6 +100,14 @@ PyFrame_GetBack(PyFrameObject *frame)
 }
 #endif
 
+static inline PyFrameObject*
+_PyFrame_GetBackBorrow(PyFrameObject *frame)
+{
+    PyFrameObject *back = PyFrame_GetBack(frame);
+    Py_XDECREF(back);
+    return back;  // borrowed reference
+}
+
 
 #if PY_VERSION_HEX < 0x030900A5
 static inline PyInterpreterState *
@@ -130,6 +130,14 @@ PyThreadState_GetFrame(PyThreadState *tstate)
 }
 #endif
 
+static inline PyFrameObject*
+_PyThreadState_GetFrameBorrow(PyThreadState *tstate)
+{
+    PyFrameObject *frame = PyThreadState_GetFrame(tstate);
+    Py_XDECREF(frame);
+    return frame;  // borrowed reference
+}
+
 
 #if PY_VERSION_HEX < 0x030900A5
 static inline PyInterpreterState *
diff --git a/tests/test_pythoncapi_compat.c b/tests/test_pythoncapi_compat.c
@@ -39,18 +39,6 @@ test_object(PyObject *self, PyObject *ignored)
     int is_type = Py_IS_TYPE(obj, Py_TYPE(obj));
     assert(is_type);
 
-    // test _Py_Borrow()
-    Py_INCREF(obj);
-    PyObject *borrowed = _Py_Borrow(obj);
-    assert(borrowed == obj);
-    assert(Py_REFCNT(obj) == refcnt);
-
-    // test _Py_XBorrow()
-    Py_INCREF(obj);
-    PyObject *xborrowed = _Py_XBorrow(obj);
-    assert(xborrowed == obj);
-    assert(Py_REFCNT(obj) == refcnt);
-
     Py_DECREF(obj);
     Py_RETURN_NONE;
 }
@@ -67,17 +55,41 @@ test_frame(PyObject *self, PyObject *ignored)
         return ASSERT_FAILED("PyThreadState_GetFrame failed");
     }
 
+    // test _PyThreadState_GetFrameBorrow()
+    Py_ssize_t frame_refcnt = Py_REFCNT(frame);
+    PyFrameObject *frame2 = _PyThreadState_GetFrameBorrow(tstate);
+    assert(frame2 == frame);
+    assert(Py_REFCNT(frame) == frame_refcnt);
+
     // test PyFrame_GetCode()
     PyCodeObject *code = PyFrame_GetCode(frame);
     assert(code != NULL);
     assert(PyCode_Check(code));
+
+    // test _PyFrame_GetCodeBorrow()
+    Py_ssize_t code_refcnt = Py_REFCNT(code);
+    PyCodeObject *code2 = _PyFrame_GetCodeBorrow(frame);
+    assert(code2 == code);
+    assert(Py_REFCNT(code) == code_refcnt);
     Py_DECREF(code);
 
     // PyFrame_GetBack()
     PyFrameObject* back = PyFrame_GetBack(frame);
     if (back != NULL) {
         assert(PyFrame_Check(back));
     }
+
+    // test _PyFrame_GetBackBorrow()
+    if (back != NULL) {
+        Py_ssize_t back_refcnt = Py_REFCNT(back);
+        PyCodeObject *back2 = _PyFrame_GetBackBorrow(frame);
+        assert(back2 == back);
+        assert(Py_REFCNT(back) == back_refcnt);
+    }
+    else {
+        PyCodeObject *back2 = _PyFrame_GetBackBorrow(frame);
+        assert(back2 == back);
+    }
     Py_XDECREF(back);
 
     Py_DECREF(frame);
