gh-103091: Add PyUnstable_Type_AssignVersionTag (#103095)

6 files changed, 49 insertions, 0 deletions

diff --git a/Doc/c-api/type.rst b/Doc/c-api/type.rst
index 7b5d1fa..69b1529 100644
--- a/Doc/c-api/type.rst
+++ b/Doc/c-api/type.rst
@@ -232,6 +232,15 @@ Type Objects
 
    .. versionadded:: 3.11
 
+.. c:function:: int PyUnstable_Type_AssignVersionTag(PyTypeObject *type)
+
+   Attempt to assign a version tag to the given type.
+
+   Returns 1 if the type already had a valid version tag or a new one was
+   assigned, or 0 if a new tag could not be assigned.
+
+   .. versionadded:: 3.12
+
 
 Creating Heap-Allocated Types
 .............................
diff --git a/Include/cpython/object.h b/Include/cpython/object.h
index 98cc51c..ce4d13c 100644
--- a/Include/cpython/object.h
+++ b/Include/cpython/object.h
@@ -564,3 +564,10 @@ PyAPI_FUNC(int) PyType_AddWatcher(PyType_WatchCallback callback);
 PyAPI_FUNC(int) PyType_ClearWatcher(int watcher_id);
 PyAPI_FUNC(int) PyType_Watch(int watcher_id, PyObject *type);
 PyAPI_FUNC(int) PyType_Unwatch(int watcher_id, PyObject *type);
+
+/* Attempt to assign a version tag to the given type.
+ *
+ * Returns 1 if the type already had a valid version tag or a new one was
+ * assigned, or 0 if a new tag could not be assigned.
+ */
+PyAPI_FUNC(int) PyUnstable_Type_AssignVersionTag(PyTypeObject *type);
diff --git a/Lib/test/test_type_cache.py b/Lib/test/test_type_cache.py
index 8502f6b..24f83cd 100644
--- a/Lib/test/test_type_cache.py
+++ b/Lib/test/test_type_cache.py
@@ -9,6 +9,7 @@ except ImportError:
 
 # Skip this test if the _testcapi module isn't available.
 type_get_version = import_helper.import_module('_testcapi').type_get_version
+type_assign_version = import_helper.import_module('_testcapi').type_assign_version
 
 
 @support.cpython_only
@@ -42,6 +43,19 @@ class TypeCacheTests(unittest.TestCase):
         self.assertEqual(len(set(all_version_tags)), 30,
                          msg=f"{all_version_tags} contains non-unique versions")
 
+    def test_type_assign_version(self):
+        class C:
+            x = 5
+
+        self.assertEqual(type_assign_version(C), 1)
+        c_ver = type_get_version(C)
+
+        C.x = 6
+        self.assertEqual(type_get_version(C), 0)
+        self.assertEqual(type_assign_version(C), 1)
+        self.assertNotEqual(type_get_version(C), 0)
+        self.assertNotEqual(type_get_version(C), c_ver)
+
 
 if __name__ == "__main__":
     support.run_unittest(TypeCacheTests)
diff --git a/Misc/NEWS.d/next/C API/2023-03-28-12-31-51.gh-issue-103091.CzZyaZ.rst b/Misc/NEWS.d/next/C API/2023-03-28-12-31-51.gh-issue-103091.CzZyaZ.rst
new file mode 100644
index 0000000..28c77b6
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2023-03-28-12-31-51.gh-issue-103091.CzZyaZ.rst
@@ -0,0 +1 @@
+Add a new C-API function to eagerly assign a version tag to a PyTypeObject: ``PyUnstable_Type_AssignVersionTag()``.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 557a6d4..30b2674 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -2733,6 +2733,18 @@ type_get_version(PyObject *self, PyObject *type)
 }
 
 
+static PyObject *
+type_assign_version(PyObject *self, PyObject *type)
+{
+    if (!PyType_Check(type)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a type");
+        return NULL;
+    }
+    int res = PyUnstable_Type_AssignVersionTag((PyTypeObject *)type);
+    return PyLong_FromLong(res);
+}
+
+
 // Test PyThreadState C API
 static PyObject *
 test_tstate_capi(PyObject *self, PyObject *Py_UNUSED(args))
@@ -3530,6 +3542,7 @@ static PyMethodDef TestMethods[] = {
     {"test_py_is_macros", test_py_is_macros, METH_NOARGS},
     {"test_py_is_funcs", test_py_is_funcs, METH_NOARGS},
     {"type_get_version", type_get_version, METH_O, PyDoc_STR("type->tp_version_tag")},
+    {"type_assign_version", type_assign_version, METH_O, PyDoc_STR("PyUnstable_Type_AssignVersionTag")},
     {"test_tstate_capi", test_tstate_capi, METH_NOARGS, NULL},
     {"frame_getlocals", frame_getlocals, METH_O, NULL},
     {"frame_getglobals", frame_getglobals, METH_O, NULL},
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 748965d..d2b77a0 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -592,6 +592,11 @@ assign_version_tag(PyTypeObject *type)
     return 1;
 }
 
+int PyUnstable_Type_AssignVersionTag(PyTypeObject *type)
+{
+    return assign_version_tag(type);
+}
+
 
 static PyMemberDef type_members[] = {
     {"__basicsize__", T_PYSSIZET, offsetof(PyTypeObject,tp_basicsize),READONLY},