Add PyModule_AddObjectRef()

vstinner · vstinner · commit c6464c69cf04 · 2021-01-11T13:07:04.000+01:00
diff --git a/README.rst b/README.rst
@@ -213,6 +213,7 @@ GC protocol::
 
 Module helper::
 
+    int PyModule_AddObjectRef(PyObject *module, const char *name, PyObject *value);
     int PyModule_AddType(PyObject *module, PyTypeObject *type);
 
 
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
@@ -32,7 +32,7 @@ extern "C" {
 
 
 // bpo-42262 added Py_NewRef() to Python 3.10.0a3
-#if PY_VERSION_HEX < 0x030a00A3 && !defined(Py_NewRef)
+#if PY_VERSION_HEX < 0x030A00A3 && !defined(Py_NewRef)
 static inline PyObject* _Py_NewRef(PyObject *obj)
 {
     Py_INCREF(obj);
@@ -43,7 +43,7 @@ static inline PyObject* _Py_NewRef(PyObject *obj)
 
 
 // bpo-42262 added Py_XNewRef() to Python 3.10.0a3
-#if PY_VERSION_HEX < 0x030a00A3 && !defined(Py_XNewRef)
+#if PY_VERSION_HEX < 0x030A00A3 && !defined(Py_XNewRef)
 static inline PyObject* _Py_XNewRef(PyObject *obj)
 {
     Py_XINCREF(obj);
@@ -246,6 +246,21 @@ PyModule_AddType(PyObject *module, PyTypeObject *type)
 #endif
 
 
+// bpo-1635741 added PyModule_AddObjectRef() to Python 3.10.0a3
+#if PY_VERSION_HEX < 0x030A00A3
+static inline int
+PyModule_AddObjectRef(PyObject *module, const char *name, PyObject *value)
+{
+    Py_XINCREF(value);
+    int res = PyModule_AddObject(module, name, value);
+    if (res < 0) {
+        Py_XDECREF(value);
+    }
+    return res;
+}
+#endif
+
+
 // bpo-40241 added PyObject_GC_IsTracked() to Python 3.9.0a6.
 // bpo-4688 added _PyObject_GC_IS_TRACKED() to Python 2.7.0a2.
 #if PY_VERSION_HEX < 0x030900A6
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
@@ -197,16 +197,10 @@ test_gc(PyObject *self, PyObject *ignored)
 }
 
 
-static PyObject *
-test_module(PyObject *self, PyObject *ignored)
+// test PyModule_AddType()
+static int
+test_module_add_type(PyObject *module)
 {
-    PyObject *module = PyImport_ImportModule("sys");
-    if (module == NULL) {
-        return NULL;
-    }
-    assert(PyModule_Check(module));
-
-    // test PyModule_AddType()
     PyTypeObject *type = &PyUnicode_Type;
 #ifdef PYTHON3
     const char *type_name = "str";
@@ -216,21 +210,69 @@ test_module(PyObject *self, PyObject *ignored)
     Py_ssize_t refcnt = Py_REFCNT(type);
 
     if (PyModule_AddType(module, type) < 0) {
-        goto error;
+        return -1;
     }
     assert(Py_REFCNT(type) == refcnt + 1);
 
     PyObject *attr = PyObject_GetAttrString(module, type_name);
     if (attr == NULL) {
-        goto error;
+        return -1;
     }
     assert(attr == (PyObject *)type);
     Py_DECREF(attr);
 
     if (PyObject_DelAttrString(module, type_name) < 0) {
-        goto error;
+        return -1;
     }
     assert(Py_REFCNT(type) == refcnt);
+    return 0;
+}
+
+
+// test PyModule_AddObjectRef()
+static int
+test_module_addobjectref(PyObject *module)
+{
+    PyObject *obj = Py_True;
+    Py_ssize_t refcnt = Py_REFCNT(obj);
+    const char *name = "test_module_addobjectref";
+
+    if (PyModule_AddObjectRef(module, name, obj) < 0) {
+        assert(Py_REFCNT(obj) == refcnt);
+        return -1;
+    }
+    assert(Py_REFCNT(obj) == refcnt + 1);
+
+    if (PyObject_DelAttrString(module, name) < 0) {
+        return -1;
+    }
+    assert(Py_REFCNT(obj) == refcnt);
+
+    // PyModule_AddObjectRef() with value=NULL must not crash
+    int res = PyModule_AddObjectRef(module, name, NULL);
+    assert(res < 0);
+    PyErr_Clear();
+
+    return 0;
+}
+
+
+static PyObject *
+test_module(PyObject *self, PyObject *ignored)
+{
+    PyObject *module = PyImport_ImportModule("sys");
+    if (module == NULL) {
+        return NULL;
+    }
+    assert(PyModule_Check(module));
+
+    if (test_module_add_type(module) < 0) {
+        goto error;
+    }
+
+    if (test_module_addobjectref(module) < 0) {
+        goto error;
+    }
 
     Py_DECREF(module);
     Py_RETURN_NONE;

Original file line number	Diff line number	Diff line change
`@@ -213,6 +213,7 @@ GC protocol::`
`213`	`213`
`214`	`214`	`Module helper::`
`215`	`215`
	`216`	`+ int PyModule_AddObjectRef(PyObject module, const char name, PyObject *value);`
`216`	`217`	`int PyModule_AddType(PyObject module, PyTypeObject type);`
`217`	`218`
`218`	`219`