do not cache sub interp capi

python · neonene · Apr 1, 2024 · Apr 1, 2024 · Apr 1, 2024 · Apr 1, 2024
commit a9dee17af5721656ecbc20c173f494567cd3bf66
@@ -194,10 +194,24 @@ typedef struct {
  * */
 #ifndef _PY_DATETIME_IMPL
 /* Define global variable for the C API and a macro for setting it. */
-static PyDateTime_CAPI *PyDateTimeAPI = NULL;
+static PyDateTime_CAPI *_pydatetimeapi_main = NULL;
 
-#define PyDateTime_IMPORT \
-    PyDateTimeAPI = (PyDateTime_CAPI *)PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0)
+static inline void _import_pydatetime(void) {
+    if (PyInterpreterState_Get() == PyInterpreterState_Main()) {
+        _pydatetimeapi_main = PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0);
+    }
+}
+#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()
 
 /* Macro for access to the UTC singleton */
 #define PyDateTime_TimeZone_UTC PyDateTimeAPI->TimeZone_UTC

@@ -2203,6 +2203,22 @@ def test_module_state_shared_in_global(self):
         subinterp_attr_id = os.read(r, 100)
         self.assertEqual(main_attr_id, subinterp_attr_id)
 
+    @unittest.skipIf(_testmultiphase is None, "test requires _testmultiphase module")
+    def test_datetime_capi(self):
+        script = textwrap.dedent("""
+            import importlib.machinery
+            import importlib.util
+            fullname = '_test_datetime_capi'
+            origin = importlib.util.find_spec('_testmultiphase').origin
+            loader = importlib.machinery.ExtensionFileLoader(fullname, origin)
+            spec = importlib.util.spec_from_loader(fullname, loader)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+        """)
+        exec(script)
+        ret = support.run_in_subinterp(script)
+        self.assertEqual(ret, 0)
+
 
 @requires_subinterpreters
 class InterpreterIDTests(unittest.TestCase):

@@ -952,3 +952,46 @@ PyInit__test_shared_gil_only(void)
 {
     return PyModuleDef_Init(&shared_gil_only_def);
 }
+
+
+#include "datetime.h"
+
+static int
+datetime_capi_exec(PyObject *m)
+{
+    _pydatetimeapi_main = NULL;
+    PyDateTime_IMPORT;
+    if (PyDateTimeAPI == NULL) {
+        return -1;
+    }
+    if (PyDateTimeAPI != PyCapsule_Import(PyDateTime_CAPSULE_NAME, 0)) {
+        return -1;
+    }
+    if (PyInterpreterState_Get() == PyInterpreterState_Main()) {
+        if (PyDateTimeAPI != _pydatetimeapi_main) {
+            return -1;
+        }
+    }
+    else {
+        if (PyDateTimeAPI == _pydatetimeapi_main) {
+            return -1;
+        }
+    }
+    return 0;
+}
+
+static PyModuleDef_Slot datetime_capi_slots[] = {
+    {Py_mod_exec, datetime_capi_exec},
+    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
+    {0, NULL},
+};
+
+static PyModuleDef datetime_capi_def = TEST_MODULE_DEF("_test_datetime_capi",
+                                                        datetime_capi_slots,
+                                                        testexport_methods);
+
+PyMODINIT_FUNC
+PyInit__test_datetime_capi(void)
+{
+    return PyModuleDef_Init(&datetime_capi_def);
+}