ENH: Add function for adding docstrings to ufuncs.

6 files changed, 133 insertions, 21 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index ce61c5f0e..f1ce48b69 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -5071,7 +5071,7 @@ add_newdoc('numpy.lib._compiled_base', 'unravel_index',
 
 add_newdoc('numpy.lib._compiled_base', 'add_docstring',
     """
-    docstring(obj, docstring)
+    add_docstring(obj, docstring)
 
     Add a docstring to a built-in obj if possible.
     If the obj already has a docstring raise a RuntimeError
@@ -5079,6 +5079,33 @@ add_newdoc('numpy.lib._compiled_base', 'add_docstring',
     raise a TypeError
     """)
 
+add_newdoc('numpy.lib._compiled_base', 'add_newdoc_ufunc',
+    """
+    add_ufunc_docstring(ufunc, new_docstring)
+
+    Replace the docstring for a ufunc with new_docstring.
+    This method will only work if the current docstring for
+    the ufunc is NULL. (At the C level, i.e. when ufunc->doc is NULL.)
+
+    Parameters
+    ----------
+    ufunc : numpy.ufunc
+        A ufunc whose current doc is NULL.
+    new_docstring : string
+        The new docstring for the ufunc.
+
+    Notes
+    -----
+
+    This method allocates memory for new_docstring on
+    the heap. Technically this creates a mempory leak, since this
+    memory will not be reclaimed until the end of the program
+    even if the ufunc itself is removed. However this will only
+    be a problem if the user is repeatedly creating ufuncs with
+    no documentation, adding documentation via add_newdoc_ufunc,
+    and then throwing away the ufunc.
+    """)
+
 add_newdoc('numpy.lib._compiled_base', 'packbits',
     """
     packbits(myarray, axis=None)
diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py
index 760617340..2e1cdd460 100644
--- a/numpy/core/code_generators/ufunc_docstrings.py
+++ b/numpy/core/code_generators/ufunc_docstrings.py
@@ -1,4 +1,13 @@
 # Docstrings for generated ufuncs
+#
+# The syntax is designed to look like the function add_newdoc is being
+# called from numpy.lib, but in this file  add_newdoc puts the docstrings
+# in a dictionary. This dictionary is used in
+# numpy/core/code_generators/generate_umath.py to generate the docstrings
+# for the ufuncs in numpy.core at the C level when the ufuncs are created
+# at compile time.
+
+
 
 docdict = {}
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 273aea996..6bee30797 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -4379,12 +4379,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
     else {
         ufunc->name = name;
     }
-    if (doc == NULL) {
-        ufunc->doc = "NULL";
-    }
-    else {
-        ufunc->doc = doc;
-    }
+    ufunc->doc = doc;
 
     /* generalized ufunc */
     ufunc->core_enabled = 0;
@@ -4839,19 +4834,35 @@ ufunc_get_doc(PyUFuncObject *ufunc)
     PyObject *outargs, *inargs, *doc;
     outargs = _makeargs(ufunc->nout, "out", 1);
     inargs = _makeargs(ufunc->nin, "x", 0);
-    if (outargs == NULL) {
-        doc = PyUString_FromFormat("%s(%s)\n\n%s",
-                                   ufunc->name,
-                                   PyString_AS_STRING(inargs),
-                                   ufunc->doc);
-    }
-    else {
-        doc = PyUString_FromFormat("%s(%s[, %s])\n\n%s",
-                                   ufunc->name,
-                                   PyString_AS_STRING(inargs),
-                                   PyString_AS_STRING(outargs),
-                                   ufunc->doc);
-        Py_DECREF(outargs);
+
+    if(ufunc->doc == NULL){
+        if(outargs == NULL){
+            doc = PyUString_FromFormat("%s(%s)\n\n",
+                                        ufunc->name,
+                                        PyString_AS_STRING(inargs));
+        }else{
+            doc = PyUString_FromFormat("%s(%s[, %s])\n\n",
+                                        ufunc->name,
+                                        PyString_AS_STRING(inargs),
+                                        PyString_AS_STRING(outargs));
+            Py_DECREF(outargs);
+        }
+    }
+    else{
+        if (outargs == NULL) {
+            doc = PyUString_FromFormat("%s(%s)\n\n%s",
+                                       ufunc->name,
+                                       PyString_AS_STRING(inargs),
+                                       ufunc->doc);
+        }
+        else {
+            doc = PyUString_FromFormat("%s(%s[, %s])\n\n%s",
+                                       ufunc->name,
+                                       PyString_AS_STRING(inargs),
+                                       PyString_AS_STRING(outargs),
+                                       ufunc->doc);
+            Py_DECREF(outargs);
+        }
     }
     Py_DECREF(inargs);
     return doc;
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index b269d98a1..f254bbacf 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -6,7 +6,7 @@ __all__ = ['select', 'piecewise', 'trim_zeros', 'copy', 'iterable',
         'histogram', 'histogramdd', 'bincount', 'digitize', 'cov', 'corrcoef',
         'msort', 'median', 'sinc', 'hamming', 'hanning', 'bartlett',
         'blackman', 'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring',
-        'meshgrid', 'delete', 'insert', 'append', 'interp']
+        'meshgrid', 'delete', 'insert', 'append', 'interp', 'add_newdoc_ufunc']
 
 import warnings
 import types
@@ -27,6 +27,7 @@ from _compiled_base import _insert, add_docstring
 from _compiled_base import digitize, bincount, interp as compiled_interp
 from arraysetops import setdiff1d
 from utils import deprecate
+from _compiled_base import add_newdoc_ufunc
 import numpy as np
 
 
@@ -3179,6 +3180,11 @@ def add_newdoc(place, obj, doc):
        (method2, docstring2), ...]
 
     This routine never raises an error.
+
+    This routine cannot modify read-only docstrings, as appear
+    in new-style classes or built-in functions. Because this
+    routine never raises an error the caller must check manually
+    that the docstrings were changed.
        """
     try:
         new = {}
diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c
index 0b6f872c3..155f29033 100644
--- a/numpy/lib/src/_compiled_base.c
+++ b/numpy/lib/src/_compiled_base.c
@@ -4,6 +4,8 @@
 #include "numpy/noprefix.h"
 #include "numpy/npy_3kcompat.h"
 #include "npy_config.h"
+#include "numpy/ufuncobject.h"
+#include "string.h"
 
 static npy_intp
 incr_slot_(double x, double *bins, npy_intp lbins)
@@ -1170,6 +1172,48 @@ arr_add_docstring(PyObject *NPY_UNUSED(dummy), PyObject *args)
 }
 
 
+/* docstring in numpy.add_newdocs.py */
+static PyObject *
+add_newdoc_ufunc(PyObject *NPY_UNUSED(dummy), PyObject *args)
+{
+    PyUFuncObject *ufunc;
+    PyObject *str;
+    char *docstr, *newdocstr;
+
+#if defined(NPY_PY3K)
+    if (!PyArg_ParseTuple(args, "O!O!", &PyUFunc_Type, &ufunc,
+                                        &PyUnicode_Type, &str)) {
+        return NULL;
+    }
+    docstr = PyBytes_AS_STRING(PyUnicode_AsUTF8String(str));
+#else
+    if (!PyArg_ParseTuple(args, "O!O!", &PyUFunc_Type, &ufunc,
+                                         &PyString_Type, &str)) {
+         return NULL;
+    }
+    docstr = PyString_AS_STRING(str);
+#endif
+
+    if (NULL != ufunc->doc) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot change docstring of ufunc with non-NULL docstring");
+        return NULL;
+    }
+
+    /*
+     * This introduces a memory leak, as the memory allocated for the doc
+     * will not be freed even if the ufunc itself is deleted. In practice
+     * this should not be a problem since the user would have to
+     * repeatedly create, document, and throw away ufuncs.
+     */
+    newdocstr = malloc(strlen(docstr) + 1);
+    strcpy(newdocstr, docstr);
+    ufunc->doc = newdocstr;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 /*  PACKBITS
  *
  *  This function packs binary (0 or 1) 1-bit per pixel arrays
@@ -1436,6 +1480,8 @@ static struct PyMethodDef methods[] = {
         METH_VARARGS | METH_KEYWORDS, NULL},
     {"add_docstring", (PyCFunction)arr_add_docstring,
         METH_VARARGS, NULL},
+    {"add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
+        METH_VARARGS, NULL},
     {"packbits", (PyCFunction)io_pack,
         METH_VARARGS | METH_KEYWORDS, NULL},
     {"unpackbits", (PyCFunction)io_unpack,
@@ -1505,6 +1551,7 @@ init_compiled_base(void)
 
     /* Import the array objects */
     import_array();
+    import_umath();
 
     /* Add some symbolic constants to the module */
     d = PyModule_GetDict(m);
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 0b6b1e19d..ba6b336ff 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1185,5 +1185,17 @@ def test_median():
     assert_allclose(np.median(a2, axis=1), [1, 4])
 
 
+class TestAdd_newdoc_ufunc(TestCase):
+
+    def test_ufunc_arg(self):
+        assert_raises(TypeError, add_newdoc_ufunc, 2, "blah")
+        assert_raises(ValueError, add_newdoc_ufunc,np.add, "blah")
+
+    def test_string_arg(self):
+        assert_raises(TypeError, add_newdoc_ufunc,np.add, 3)
+
+
+
+
 if __name__ == "__main__":
     run_module_suite()