apply PR: Access C-API via PyInterpreterState

python · neonene · Apr 26, 2024 · Apr 26, 2024 · Apr 26, 2024 · Apr 26, 2024
commit 49b16c00464c063a8a190c6eefd1721f32d78b66
@@ -186,34 +186,34 @@ typedef struct {
 } PyDateTime_CAPI;
 
 #define PyDateTime_CAPSULE_NAME "datetime.datetime_CAPI"
-
+#define PyDateTime_INTERNAL_CAPSULE_NAME "datetime.datetime_CAPI_INTERNAL"
 
 /* This block is only used as part of the public API and should not be
  * included in _datetimemodule.c, which does not use the C API capsule.
  * See bpo-35081 for more details.
  * */
 #ifndef _PY_DATETIME_IMPL
-/* Define global variable for the C API and a macro for setting it. */
-static PyDateTime_CAPI *_pydatetimeapi_main = NULL;
+static PyDateTime_CAPI *
+_PyDateTimeAPI_not_ready(void)
+{
+    return NULL;
+}
+// typedef PyDateTime_CAPI *(*datetime_api_getfunc)(void);
+// static datetime_api_getfunc _PyDateTimeAPI_Get = _PyDateTimeAPI_not_ready;
+static PyDateTime_CAPI *(*_PyDateTimeAPI_Get)(void) = _PyDateTimeAPI_not_ready;
+
 
 static inline void
-_import_pydatetime(void) {
-    if (PyInterpreterState_Get() == PyInterpreterState_Main()) {
-        _pydatetimeapi_main = PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0);
+_PyDateTimeAPI_Import(void)
+{
+    void *(*func)(void) = PyCapsule_Import(PyDateTime_INTERNAL_CAPSULE_NAME, 0);
+    if (func) {
+        _PyDateTimeAPI_Get = func();
     }
 }
-#define PyDateTime_IMPORT _import_pydatetime()
 
-static inline PyDateTime_CAPI *
-_get_pydatetime_api(void) {
-    if (PyInterpreterState_Get() == PyInterpreterState_Main()) {
-        return _pydatetimeapi_main;
-    }
-    else {
-        return PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0);
-    }
-}
-#define PyDateTimeAPI _get_pydatetime_api()
+#define PyDateTimeAPI _PyDateTimeAPI_Get()
+#define PyDateTime_IMPORT _PyDateTimeAPI_Import()
 
 /* Macro for access to the UTC singleton */
 #define PyDateTime_TimeZone_UTC PyDateTimeAPI->TimeZone_UTC

@@ -236,6 +236,7 @@ struct _is {
     // more comments.
     struct _obmalloc_state *obmalloc;
 
+    void *datetime_capi;
     PyObject *audit_hooks;
     PyType_WatchCallback type_watchers[TYPE_MAX_WATCHERS];
     PyCode_WatchCallback code_watchers[CODE_MAX_WATCHERS];

@@ -91,7 +91,7 @@ def test_name_cleanup(self):
                     if not name.startswith('__') and not name.endswith('__'))
         allowed = set(['MAXYEAR', 'MINYEAR', 'date', 'datetime',
                        'datetime_CAPI', 'time', 'timedelta', 'timezone',
-                       'tzinfo', 'UTC', 'sys'])
+                       'tzinfo', 'UTC', 'sys', 'datetime_CAPI_INTERNAL'])
         self.assertEqual(names - allowed, set([]))
 
     def test_divide_and_round(self):

@@ -2295,6 +2295,7 @@ def test_datetime_capi_client(self):
             spec.loader.exec_module(module)
         """)
         exec(script)  # run main interp first
+        exec(script)  # run main interp twice
         ret = support.run_in_subinterp(script)
         self.assertEqual(ret, 0)
 

@@ -108,6 +108,30 @@ find_state_left_or_right(PyObject *left, PyObject *right)
 
 #define find_module_state_by_def(obj) find_module_state_by_def((PyTypeObject *)(obj))
 
+static inline void
+set_datetime_capi_by_interp(PyDateTime_CAPI *capi)
+{
+    _PyInterpreterState_GET()->datetime_capi = capi;
+}
+
+static PyDateTime_CAPI *
+_PyDateTimeAPI_Get(void)
+{
+    return (PyDateTime_CAPI *)_PyInterpreterState_GET()->datetime_capi;
+}
+
+static void *
+_PyDateTimeAPI_Import(void)
+{
+    PyDateTime_CAPI *capi = PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0);
+    if (capi) {
+        // PyInit__datetime() is not called when the module is already loaded
+        // with single-phase init.
+        set_datetime_capi_by_interp(capi);
+        return (void *)_PyDateTimeAPI_Get;
+    }
+    return NULL;
+}
 
 /* We require that C int be at least 32 bits, and use int virtually
  * everywhere.  In just a few cases we use a temp long, where a Python
@@ -7457,6 +7481,20 @@ _datetime_exec(PyObject *module)
         goto error;
     }
 
+    capsule = PyCapsule_New(_PyDateTimeAPI_Import,
+                            PyDateTime_INTERNAL_CAPSULE_NAME, NULL);
+    if (capsule == NULL) {
+        PyMem_Free(capi);
+        goto error;
+    }
+    if (PyModule_Add(module, "datetime_CAPI_INTERNAL", capsule) < 0) {
+        PyMem_Free(capi);
+        goto error;
+    }
+
+    /* Ensure that the newest capi is used on multi-phase init */
+    set_datetime_capi_by_interp(capi);
+
     /* A 4-year cycle has an extra leap day over what we'd get from
      * pasting together 4 single years.
      */

diff --git a/Modules/_testcapi/datetime.c b/Modules/_testcapi/datetime.c
@@ -3,11 +3,23 @@
 #include "datetime.h"             // PyDateTimeAPI
 
 
+static int test_run_counter = 0;
+
 static PyObject *
 test_datetime_capi(PyObject *self, PyObject *args)
 {
-    // PyDateTimeAPI cannot be carried over
-    // with multi-phase init enabled.
+    if (PyDateTimeAPI) {
+        if (test_run_counter) {
+            /* Probably regrtest.py -R */
+            Py_RETURN_NONE;
+        }
+        else {
+            PyErr_SetString(PyExc_AssertionError,
+                            "PyDateTime_CAPI somehow initialized");
+            return NULL;
+        }
+    }
+    test_run_counter++;
     PyDateTime_IMPORT;
 
     if (PyDateTimeAPI) {

@@ -957,44 +957,62 @@ PyInit__test_shared_gil_only(void)
 #include "datetime.h"
 
 static int
-datetime_capi_client_exec(PyObject *m)
+datetime_capi_import_with_error(void)
 {
+    static int is_datetime_multiphase = -1;
     int ismain = PyInterpreterState_Get() == PyInterpreterState_Main();
-    if (ismain) {
-        _pydatetimeapi_main = NULL;
+    if (ismain && is_datetime_multiphase < 0) {
+        PyObject *module = PyImport_ImportModule("_datetime");
+        if (module == NULL) {
+            return -1;
+        }
+        PyModuleDef *def = PyModule_GetDef(module);
+        Py_DECREF(module);
+        if (def && def->m_size >= 0) {
+            is_datetime_multiphase = 1;
+        }
+        else {
+            is_datetime_multiphase = 0;
+        }
+    }
+    if (is_datetime_multiphase < 0) {
+        PyErr_SetString(PyExc_AssertionError,
+                        "Main interpreter must be loaded first.");
+        return -1;
     }
 
-    PyDateTime_IMPORT;
-    PyErr_Clear();
-    PyDateTime_CAPI *capi = PyDateTimeAPI;
-    PyErr_Clear();
-    if (capi != PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0)) {
+    _PyDateTimeAPI_Import();
+    if (!PyErr_Occurred()) {
+        return 0;
+    }
+#ifdef Py_GIL_DISABLED
+    if (!ismain && !is_datetime_multiphase) {
+        // _datetime module and Capsule are not imported
+        PyErr_WriteUnraisable(NULL);
+        return 0;
+    }
+#endif
+    return -1;
+}
+
+static int
+datetime_capi_client_exec(PyObject *m)
+{
+    _PyDateTimeAPI_Get = _PyDateTimeAPI_not_ready;
+    if (_PyDateTimeAPI_Get() != NULL) {
+        PyErr_SetString(PyExc_AssertionError,
+                        "DateTime API is expected to remain NULL.");
         return -1;
     }
-    PyErr_Clear();
-    if (ismain) {
-        if (capi != _pydatetimeapi_main) {
-            return -1;
-        }
+    if (datetime_capi_import_with_error() < 0) {
+        return -1;
     }
-    else {
-        if (capi == _pydatetimeapi_main) {
-            PyObject *module = PyImport_ImportModule("_datetime");
-            if (module == NULL) {
-                return -1;
-            }
-            PyModuleDef *def = PyModule_GetDef(module);
-            Py_DECREF(module);
-            if (def) {
-                // multi-phase init
-                return -1;
-            }
-            else {
-                // legacy init (shared module)
-                return 0;
-            }
-        }
+    if (PyDateTimeAPI != PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0)) {
+        PyErr_SetString(PyExc_AssertionError,
+                        "DateTime API does not match Capsule CAPI.");
+        return -1;
     }
+    PyErr_Clear();
     return 0;
 }
 

@@ -437,7 +437,7 @@ Modules/_tkinter.c	-	trbInCmd	-
 ## initialized once
 
 ## other
-Include/datetime.h	-	_pydatetimeapi_main	-
+Include/datetime.h	-	_PyDateTimeAPI_Get	-
 Modules/_ctypes/cfield.c	_ctypes_get_fielddesc	initialized	-
 Modules/_ctypes/malloc_closure.c	-	_pagesize	-
 Modules/_cursesmodule.c	-	initialised	-

@@ -591,6 +591,7 @@ Modules/_testmultiphase.c	-	slots_exec_unreported_exception	-
 Modules/_testmultiphase.c	-	slots_nonmodule_with_exec_slots	-
 Modules/_testmultiphase.c	-	testexport_methods	-
 Modules/_testmultiphase.c	-	uninitialized_def	-
+Modules/_testmultiphase.c	datetime_capi_import_with_error	is_datetime_multiphase	-
 Modules/_testsinglephase.c	-	global_state	-
 Modules/_xxtestfuzz/_xxtestfuzz.c	-	_fuzzmodule	-
 Modules/_xxtestfuzz/_xxtestfuzz.c	-	module_methods	-