ENH: Add a __dict__ to ufunc objects and use it to allow overriding __doc__

mtsokol · ngoldbaum · mtsokol · commit d6bafc61c088 · 2024-11-11T18:41:25.000Z
Co-authored-by: Nathan Goldbaum &lt;nathan.goldbaum@gmail.com&gt;
diff --git a/numpy/_core/include/numpy/ufuncobject.h b/numpy/_core/include/numpy/ufuncobject.h
@@ -170,8 +170,10 @@ typedef struct _tagPyUFuncObject {
          * with the dtypes for the inputs and outputs.
          */
         PyUFunc_TypeResolutionFunc *type_resolver;
-        /* Was the legacy loop resolver */
-        void *reserved2;
+
+        /* A dictionary to monkeypatch ufuncs */
+        PyObject *dict;
+
         /*
          * This was blocked off to be the "new" inner loop selector in 1.7,
          * but this was never implemented. (This is also why the above
diff --git a/numpy/_core/src/multiarray/npy_static_data.c b/numpy/_core/src/multiarray/npy_static_data.c
@@ -63,6 +63,7 @@ intern_strings(void)
     INTERN_STRING(__dlpack__, "__dlpack__");
     INTERN_STRING(pyvals_name, "UFUNC_PYVALS_NAME");
     INTERN_STRING(legacy, "legacy");
+    INTERN_STRING(__doc__, "__doc__");
     return 0;
 }
 
diff --git a/numpy/_core/src/multiarray/npy_static_data.h b/numpy/_core/src/multiarray/npy_static_data.h
@@ -38,6 +38,7 @@ typedef struct npy_interned_str_struct {
     PyObject *__dlpack__;
     PyObject *pyvals_name;
     PyObject *legacy;
+    PyObject *__doc__;
 } npy_interned_str_struct;
 
 /*
diff --git a/numpy/_core/src/umath/ufunc_object.c b/numpy/_core/src/umath/ufunc_object.c
@@ -4771,6 +4771,7 @@ PyUFunc_FromFuncAndDataAndSignatureAndIdentity(PyUFuncGenericFunction *func, voi
             return NULL;
         }
     }
+    ufunc->dict = PyDict_New();
     /*
      * TODO: I tried adding a default promoter here (either all object for
      *       some special cases, or all homogeneous).  Those are reasonable
@@ -6411,6 +6412,15 @@ ufunc_get_doc(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
 {
     PyObject *doc;
 
+    // If there is a __doc__ in the instance __dict__, use it.
+    int result = PyDict_GetItemRef(ufunc->dict, npy_interned_str.__doc__, &doc);
+    if (result == -1) {
+        return NULL;
+    }
+    else if (result == 1) {
+        return doc;
+    }
+
     if (npy_cache_import_runtime(
             "numpy._core._internal", "_ufunc_doc_signature_formatter",
             &npy_runtime_imports._ufunc_doc_signature_formatter) == -1) {
@@ -6434,6 +6444,24 @@ ufunc_get_doc(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
     return doc;
 }
 
+static int
+ufunc_set_doc(PyUFuncObject *ufunc, PyObject *doc, void *NPY_UNUSED(ignored))
+{
+    if (doc == NULL) {
+        int result = PyDict_Contains(ufunc->dict, npy_interned_str.__doc__);
+        if (result == -1) {
+            return NULL;
+        }
+        else if (result == 1) {
+            return PyDict_DelItem(ufunc->dict, npy_interned_str.__doc__);
+        }
+        else {
+            return 0;
+        }
+    } else {
+        return PyDict_SetItem(ufunc->dict, npy_interned_str.__doc__, doc);
+    }
+}
 
 static PyObject *
 ufunc_get_nin(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
@@ -6519,8 +6547,8 @@ ufunc_get_signature(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
 
 static PyGetSetDef ufunc_getset[] = {
     {"__doc__",
-        (getter)ufunc_get_doc,
-        NULL, NULL, NULL},
+        (getter)ufunc_get_doc, (setter)ufunc_set_doc,
+        NULL, NULL},
     {"nin",
         (getter)ufunc_get_nin,
         NULL, NULL, NULL},
@@ -6549,6 +6577,17 @@ static PyGetSetDef ufunc_getset[] = {
 };
 
 
+/******************************************************************************
+ ***                          UFUNC MEMBERS                                 ***
+ *****************************************************************************/
+
+static PyMemberDef ufunc_members[] = {
+    {"__dict__", T_OBJECT, offsetof(PyUFuncObject, dict),
+     READONLY},
+    {NULL},
+};
+
+
 /******************************************************************************
  ***                        UFUNC TYPE OBJECT                               ***
  *****************************************************************************/
@@ -6568,6 +6607,12 @@ NPY_NO_EXPORT PyTypeObject PyUFunc_Type = {
     .tp_traverse = (traverseproc)ufunc_traverse,
     .tp_methods = ufunc_methods,
     .tp_getset = ufunc_getset,
+    .tp_getattro = PyObject_GenericGetAttr,
+    .tp_setattro = PyObject_GenericSetAttr,
+    // TODO when Python 3.12 is the minimum supported version,
+    // use Py_TPFLAGS_MANAGED_DICT
+    .tp_members = ufunc_members,
+    .tp_dictoffset = offsetof(PyUFuncObject, dict),
 };
 
 /* End of code for ufunc objects */
diff --git a/numpy/_core/tests/test_umath.py b/numpy/_core/tests/test_umath.py
@@ -4016,6 +4016,26 @@ def test_array_ufunc_direct_call(self):
         res = a.__array_ufunc__(np.add, "__call__", a, a)
         assert_array_equal(res, a + a)
 
+    def test_ufunc_docstring(self):
+        original_doc = np.add.__doc__
+        new_doc = "new docs"
+
+        np.add.__doc__ = new_doc
+        assert np.add.__doc__ == new_doc
+        assert np.add.__dict__["__doc__"] == new_doc
+
+        del np.add.__doc__
+        assert np.add.__doc__ == original_doc
+        assert np.add.__dict__ == {}
+
+        np.add.__dict__["other"] = 1
+        np.add.__dict__["__doc__"] = new_doc
+        assert np.add.__doc__ == new_doc
+
+        del np.add.__dict__["__doc__"]
+        assert np.add.__doc__ == original_doc
+
+
 class TestChoose:
     def test_mixed(self):
         c = np.array([True, True])

Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,7 @@ intern_strings(void)`
`63`	`63`	`INTERN_STRING(__dlpack__, "__dlpack__");`
`64`	`64`	`INTERN_STRING(pyvals_name, "UFUNC_PYVALS_NAME");`
`65`	`65`	`INTERN_STRING(legacy, "legacy");`
	`66`	`+ INTERN_STRING(__doc__, "__doc__");`
`66`	`67`	`return 0;`
`67`	`68`	`}`
`68`	`69`