MNT: move default extobj contextvar to global data dict

Author: ngoldbaum

numpy/_core/src/multiarray/multiarraymodule.h

@@ -38,6 +38,20 @@ typedef struct npy_ma_global_data_struct {
      */
     PyObject *default_truediv_type_tup;
 
+    /*
+     * Used to set up the default extobj context variable
+     *
+     * This is immutable and set at module initialization so can be used
+     * without acquiring the global data mutex
+     */
+    PyObject *default_extobj_capsule;
+
+    /*
+     * The global ContextVar to store the extobject. It is exposed to Python
+     * as `_extobj_contextvar`.
+     */
+    PyObject *npy_extobj_contextvar;
+
     /*
      * A reference to ndarray's implementations for __array_*__ special methods
      */

numpy/_core/src/umath/extobj.c

@@ -18,14 +18,6 @@
 #include "common.h"
 
 
-/*
- * The global ContextVar to store the extobject. It is exposed to Python
- * as `_extobj_contextvar`.
- */
-static PyObject *default_extobj_capsule = NULL;
-NPY_NO_EXPORT PyObject *npy_extobj_contextvar = NULL;
-
-
 #define UFUNC_ERR_IGNORE 0
 #define UFUNC_ERR_WARN   1
 #define UFUNC_ERR_RAISE  2
@@ -130,7 +122,8 @@ fetch_curr_extobj_state(npy_extobj *extobj)
 {
     PyObject *capsule;
     if (PyContextVar_Get(
-            npy_extobj_contextvar, default_extobj_capsule, &capsule) < 0) {
+            npy_ma_global_data->npy_extobj_contextvar,
+            npy_ma_global_data->default_extobj_capsule, &capsule) < 0) {
         return -1;
     }
     npy_extobj *obj = PyCapsule_GetPointer(capsule, "numpy.ufunc.extobj");
@@ -164,15 +157,15 @@ init_extobj(void)
         }
     }
 
-    default_extobj_capsule = make_extobj_capsule(
+    npy_ma_global_data->default_extobj_capsule = make_extobj_capsule(
             NPY_BUFSIZE, UFUNC_ERR_DEFAULT, Py_None);
-    if (default_extobj_capsule == NULL) {
+    if (npy_ma_global_data->default_extobj_capsule == NULL) {
         return -1;
     }
-    npy_extobj_contextvar = PyContextVar_New(
-            "numpy.ufunc.extobj", default_extobj_capsule);
-    if (npy_extobj_contextvar == NULL) {
-        Py_CLEAR(default_extobj_capsule);
+    npy_ma_global_data->npy_extobj_contextvar = PyContextVar_New(
+            "numpy.ufunc.extobj", npy_ma_global_data->default_extobj_capsule);
+    if (npy_ma_global_data->npy_extobj_contextvar == NULL) {
+        Py_CLEAR(npy_ma_global_data->default_extobj_capsule);
         return -1;
     }
     return 0;
@@ -213,7 +206,7 @@ errmodeconverter(PyObject *obj, int *mode)
 /*
  * This function is currently exposed as `umath._seterrobj()`, it is private
  * and returns a capsule representing the errstate.  This capsule is then
- * assigned to the `npy_extobj_contextvar` in Python.
+ * assigned to the `_extobj_contextvar` in Python.
  */
 NPY_NO_EXPORT PyObject *
 extobj_make_extobj(PyObject *NPY_UNUSED(mod),

numpy/_core/src/umath/extobj.h

@@ -4,9 +4,6 @@
 #include <numpy/ndarraytypes.h>  /* for NPY_NO_EXPORT */
 
 
-/* For the private exposure of the extobject contextvar to Python */
-extern NPY_NO_EXPORT PyObject *npy_extobj_contextvar;
-
 /*
  * Represent the current ufunc error (and buffer) state.  we are using a
  * capsule for now to store this, but it could make sense to refactor it into

numpy/_core/src/umath/umathmodule.c

@@ -273,8 +273,8 @@ int initumath(PyObject *m)
 #undef ADDSCONST
     PyModule_AddIntConstant(m, "UFUNC_BUFSIZE_DEFAULT", (long)NPY_BUFSIZE);
 
-    Py_INCREF(npy_extobj_contextvar);
-    PyModule_AddObject(m, "_extobj_contextvar", npy_extobj_contextvar);
+    Py_INCREF(npy_ma_global_data->npy_extobj_contextvar);
+    PyModule_AddObject(m, "_extobj_contextvar", npy_ma_global_data->npy_extobj_contextvar);
 
     PyModule_AddObject(m, "PINF", PyFloat_FromDouble(NPY_INFINITY));
     PyModule_AddObject(m, "NINF", PyFloat_FromDouble(-NPY_INFINITY));