Merge pull request numpy#27735 from mtsokol/ufunc-object-dict

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

Author: ngoldbaum

doc/release/upcoming_changes/27735.deprecation.rst

@@ -0,0 +1,2 @@
+* ``_add_newdoc_ufunc`` is now deprecated. ``ufunc.__doc__ = newdoc`` should
+  be used instead.

doc/release/upcoming_changes/27735.new_feature.rst

@@ -0,0 +1,4 @@
+* UFuncs now support ``__dict__`` attribute and allow overriding ``__doc__``
+  (either directly or via ``ufunc.__dict__["__doc__"]``). ``__dict__`` can be
+  used to also override other properties, such as ``__module__`` or
+  ``__qualname__``.

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

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;
 }
 

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;
 
 /*

numpy/_core/src/umath/_struct_ufunc_tests.c

@@ -133,8 +133,8 @@ PyMODINIT_FUNC PyInit__struct_ufunc_tests(void)
     import_umath();
 
     add_triplet = PyUFunc_FromFuncAndData(NULL, NULL, NULL, 0, 2, 1,
-                                    PyUFunc_None, "add_triplet",
-                                    "add_triplet_docstring", 0);
+                                          PyUFunc_None, "add_triplet",
+                                          NULL, 0);
 
     dtype_dict = Py_BuildValue("[(s, s), (s, s), (s, s)]",
                                "f0", "u8", "f1", "u8", "f2", "u8");

numpy/_core/src/umath/ufunc_object.c

@@ -4687,6 +4687,7 @@ PyUFunc_FromFuncAndDataAndSignatureAndIdentity(PyUFuncGenericFunction *func, voi
     ufunc->core_signature = NULL;
     ufunc->core_enabled = 0;
     ufunc->obj = NULL;
+    ufunc->dict = NULL;
     ufunc->core_num_dims = NULL;
     ufunc->core_num_dim_ix = 0;
     ufunc->core_offsets = NULL;
@@ -4771,6 +4772,11 @@ PyUFunc_FromFuncAndDataAndSignatureAndIdentity(PyUFuncGenericFunction *func, voi
             return NULL;
         }
     }
+    ufunc->dict = PyDict_New();
+    if (ufunc->dict == NULL) {
+        Py_DECREF(ufunc);
+        return NULL;
+    }
     /*
      * TODO: I tried adding a default promoter here (either all object for
      *       some special cases, or all homogeneous).  Those are reasonable
@@ -5177,6 +5183,7 @@ ufunc_dealloc(PyUFuncObject *ufunc)
         Py_DECREF(ufunc->identity_value);
     }
     Py_XDECREF(ufunc->obj);
+    Py_XDECREF(ufunc->dict);
     Py_XDECREF(ufunc->_loops);
     if (ufunc->_dispatch_cache != NULL) {
         PyArrayIdentityHash_Dealloc(ufunc->_dispatch_cache);
@@ -5197,6 +5204,7 @@ ufunc_traverse(PyUFuncObject *self, visitproc visit, void *arg)
     if (self->identity == PyUFunc_IdentityValue) {
         Py_VISIT(self->identity_value);
     }
+    Py_VISIT(self->dict);
     return 0;
 }
 
@@ -6411,6 +6419,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 +6451,15 @@ 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) {
+        return PyDict_DelItem(ufunc->dict, npy_interned_str.__doc__);
+    } else {
+        return PyDict_SetItem(ufunc->dict, npy_interned_str.__doc__, doc);
+    }
+}
 
 static PyObject *
 ufunc_get_nin(PyUFuncObject *ufunc, void *NPY_UNUSED(ignored))
@@ -6519,8 +6545,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 +6575,17 @@ static PyGetSetDef ufunc_getset[] = {
 };
 
 
+/******************************************************************************
+ ***                          UFUNC MEMBERS                                 ***
+ *****************************************************************************/
+
+static PyMemberDef ufunc_members[] = {
+    {"__dict__", T_OBJECT, offsetof(PyUFuncObject, dict),
+     READONLY},
+    {NULL},
+};
+
+
 /******************************************************************************
  ***                        UFUNC TYPE OBJECT                               ***
  *****************************************************************************/
@@ -6568,6 +6605,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 */

numpy/_core/src/umath/umathmodule.c

@@ -167,6 +167,13 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) {
 PyObject *
 add_newdoc_ufunc(PyObject *NPY_UNUSED(dummy), PyObject *args)
 {
+
+    /* 2024-11-12, NumPy 2.2 */
+    if (DEPRECATE("_add_newdoc_ufunc is deprecated. "
+                  "Use `ufunc.__doc__ = newdoc` instead.") < 0) {
+        return NULL;
+    }
+
     PyUFuncObject *ufunc;
     PyObject *str;
     if (!PyArg_ParseTuple(args, "O!O!:_add_newdoc_ufunc", &PyUFunc_Type, &ufunc,

numpy/_core/tests/test_deprecations.py

@@ -15,6 +15,7 @@
     )
 
 from numpy._core._multiarray_tests import fromstring_null_term_c_api
+import numpy._core._struct_ufunc_tests as struct_ufunc
 
 try:
     import pytz
@@ -732,3 +733,13 @@ def test_deprecated(self):
                 self.assert_deprecated(np.save, args=sample_args,
                                     kwargs={'allow_pickle': allow_pickle,
                                             'fix_imports': False})
+
+
+class TestAddNewdocUFunc(_DeprecationTestCase):
+    # Deprecated in Numpy 2.2, 2024-11
+    def test_deprecated(self):
+        self.assert_deprecated(
+            lambda: np._core.umath._add_newdoc_ufunc(
+                struct_ufunc.add_triplet, "new docs"
+            )
+        )

numpy/_core/tests/test_umath.py

@@ -4016,6 +4016,28 @@ 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
+        del np.add.__dict__["other"]
+        assert np.add.__dict__ == {}
+
+
 class TestChoose:
     def test_mixed(self):
         c = np.array([True, True])
@@ -4862,9 +4884,11 @@ def func():
 
 
 class TestAdd_newdoc_ufunc:
+    @pytest.mark.filterwarnings("ignore:_add_newdoc_ufunc:DeprecationWarning")
     def test_ufunc_arg(self):
         assert_raises(TypeError, ncu._add_newdoc_ufunc, 2, "blah")
         assert_raises(ValueError, ncu._add_newdoc_ufunc, np.add, "blah")
 
+    @pytest.mark.filterwarnings("ignore:_add_newdoc_ufunc:DeprecationWarning")
     def test_string_arg(self):
         assert_raises(TypeError, ncu._add_newdoc_ufunc, np.add, 3)