From 180750b9a4a4dc46be32ed6d738329b2ab97abfa Mon Sep 17 00:00:00 2001
From: Chris Jordan-Squire <cjordan1@uw.edu>
Date: Fri, 5 Aug 2011 10:43:51 -0500
Subject: ENH: Add function for adding docstrings to ufuncs.

---
 numpy/lib/function_base.py            |  8 +++++-
 numpy/lib/src/_compiled_base.c        | 47 +++++++++++++++++++++++++++++++++++
 numpy/lib/tests/test_function_base.py | 12 +++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)

(limited to 'numpy/lib')

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()
-- 
cgit v1.2.1