Merge pull request #9026 from eric-wieser/ufunc_docstrings

ENH: Show full PEP 457 argument lists for ufuncs

1 files changed, 45 insertions, 0 deletions

diff --git a/numpy/core/_internal.py b/numpy/core/_internal.py
index f2f24bcc5..c890eba17 100644
--- a/numpy/core/_internal.py
+++ b/numpy/core/_internal.py
@@ -670,3 +670,48 @@ def array_ufunc_errmsg_formatter(ufunc, method, *inputs, **kwargs):
     return ('operand type(s) do not implement __array_ufunc__'
             '({!r}, {!r}, {}): {}'
             .format(ufunc, method, args_string, types_string))
+
+def _ufunc_doc_signature_formatter(ufunc):
+    """
+    Builds a signature string which resembles PEP 457
+
+    This is used to construct the first line of the docstring
+    """
+
+    # input arguments are simple
+    if ufunc.nin == 1:
+        in_args = 'x'
+    else:
+        in_args = ', '.join('x{}'.format(i+1) for i in range(ufunc.nin))
+
+    # output arguments are both keyword or positional
+    if ufunc.nout == 0:
+        out_args = ', /, out=()'
+    elif ufunc.nout == 1:
+        out_args = ', /, out=None'
+    else:
+        out_args = '[, {positional}], / [, out={default}]'.format(
+            positional=', '.join(
+                'out{}'.format(i+1) for i in range(ufunc.nout)),
+            default=repr((None,)*ufunc.nout)
+        )
+
+    # keyword only args depend on whether this is a gufunc
+    kwargs = (
+        ", casting='same_kind'"
+        ", order='K'"
+        ", dtype=None"
+        ", subok=True"
+        "[, signature"
+        ", extobj]"
+    )
+    if ufunc.signature is None:
+        kwargs = ", where=True" + kwargs
+
+    # join all the parts together
+    return '{name}({in_args}{out_args}, *{kwargs})'.format(
+        name=ufunc.__name__,
+        in_args=in_args,
+        out_args=out_args,
+        kwargs=kwargs
+    )