numpy · chrisjordansquire · Aug 5, 2011 · Aug 5, 2011 · Aug 15, 2011 · charris
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -5071,14 +5071,41 @@ def luf(lamdaexpr, *args, **kwargs):
 
 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
     If this routine does not know how to add a docstring to the object
     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
@@ -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
@@ -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
@@ -6,7 +6,7 @@
         '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 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
@@ -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,49 @@ 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 (!PyArg_ParseTuple(args, "O!O!", &PyUFunc_Type, &ufunc,
+                                        &PyString_Type, &str)) {
+        return NULL;
+    }
+
+    if(NULL != ufunc->doc){
+        PyErr_SetString(PyExc_ValueError, 
+                        "Cannot change docstring of ufunc "
+                        "with non-NULL docstring");
+        return NULL;
+    }
+
+#if defined(NPY_PY3K)
+    docstr = PyBytes_AS_STRING(PyUnicode_AsUTF8String(str));
+#else
+    docstr = PyString_AS_STRING(str);
+#endif
+
+    /* 
+     * 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 +1481,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 +1552,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
@@ -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()