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

… `__doc__` Co-authored-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>
numpy · ngoldbaum · Nov 13, 2024 · Nov 11, 2024 · Nov 12, 2024 · Nov 12, 2024
commit 3eafa3f34a8e1a992e2ae365d134a45a473fd3ad
diff --git a/doc/release/upcoming_changes/27735.new_feature.rst b/doc/release/upcoming_changes/27735.new_feature.rst
@@ -0,0 +1,2 @@
+* UFuncs new support `__dict__` attribute and allow overriding
+  `__doc__` (either directly or via `ufunc.__dict__["__doc__"]`).
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,20 @@ 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 PyDict_DelItem(ufunc->dict, npy_interned_str.__doc__);
+        } else {
+            return result;
+        }
+    } else {
+        return PyDict_SetItem(ufunc->dict, npy_interned_str.__doc__, doc);
+    }
+}

 static PyObject *
 ufunc_get_nin(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
@@ -6519,8 +6543,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 +6573,17 @@ static PyGetSetDef ufunc_getset[] = {
 };


+/******************************************************************************
+ ***                          UFUNC MEMBERS                                 ***
+ *****************************************************************************/
+
+static PyMemberDef ufunc_members[] = {
+    {"__dict__", T_OBJECT, offsetof(PyUFuncObject, dict),
+     READONLY},
+    {NULL},
+};
+
+
 /******************************************************************************
 ***                        UFUNC TYPE OBJECT                               ***
 *****************************************************************************/
@@ -6568,6 +6603,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 change
@@ -0,0 +1,2 @@
+    * UFuncs new support `__dict__` attribute and allow overriding
+      `__doc__` (either directly or via `ufunc.__dict__["__doc__"]`).
                   rgommers marked this conversation as resolved.
              

          
            Show resolved
            Hide resolved
        

                
  
            
  

  
    
  
    
      

              Uh oh!

              
There was an error while loading. Please reload this page.