Repository: numpy/numpy
Pull request: ENH: Enabled the use of numpy.vectorize as a decorator
<!--         ----------------------------------------------------------------
                MAKE SURE YOUR PR GETS THE ATTENTION IT DESERVES!
                ----------------------------------------------------------------

*  FORMAT IT RIGHT:
      https://www.numpy.org/devdocs/dev/development_workflow.html#writing-the-commit-message

*  IF IT'S A NEW FEATURE OR API CHANGE, TEST THE WATERS:
      https://www.numpy.org/devdocs/dev/development_workflow.html#get-the-mailing-list-s-opinion

*  HIT ALL THE GUIDELINES:
      https://numpy.org/devdocs/dev/index.html#guidelines

*  WHAT TO DO IF WE HAVEN'T GOTTEN BACK TO YOU:
      https://www.numpy.org/devdocs/dev/development_workflow.html#getting-your-pr-reviewed
-->
This closes #9477 and #23021, as part of a small project suggested by @seberg.

Problem statement:
vectorize-style decorator
Would be really handy to have a `vectorize`-like implementation that acts like a Python decorator. To be clear, `vectorize` almost satisfies that need, but one cannot use it as a decorator and pass keyword arguments. There are certainly easy ways around this, but it would be nice to have an accepted way to do this in NumPy for simple reuse.

Hints from issue discussion:
Can you give an example of proposed syntax?
I guess some syntax like this would be good (name is arbitrary here).

```python
@dec_vectorize(otypes=None, doc=None, excluded=None, cache=False, signature=None)
def my_func(arr):
   ...
```

This means that `dec_vectorize` (or otherwise named) would look something like this.

```python
def dec_vectorize(otypes=None, doc=None, excluded=None, cache=False, signature=None):
    def wrapper(pyfunc):
        return np.vectorize(pyfunc,
                            otypes=otypes,
                            doc=doc,
                            excluded=excluded,
                            cache=cache,
                            signature=signature)
    return wrapper
```

Though this is rough and certainly open to changes.
Can't we just make the func argument to `vectorize` optional, and have it return a decorator if its ommited?
Completely agree. That sounds much better actually.
You can already do this with `functools.partial`:
```python
@partial(np.vectorize, otypes=None, doc=None, excluded=None, cache=False, signature=None)
def my_func(arr):
   ...
```

I think this is a pretty reasonable solution.

> Can't we just make the func argument to vectorize optional, and have it return a decorator if its ommited?

`vectorize` is a class, so it would require jumping through some hoops (using `__new__`, I think) to make this work.
This is doable without any magical hacks, with a small amount of added complexity on the `vectorize` class.  (basically, if a function hasn't been provided, then `__call__` will perform the final bits of initialization and return `self`)

```python
>>> @np.vectorize
... def identity(x):
...     return x
...
>>> @np.vectorize(otypes=[float])
... def as_float(x):
...     return x
...
>>> identity([0, 1, 2, 3])
array([0, 1, 2, 3])
>>> as_float([0, 1, 2, 3])
array([0., 1., 2., 3.])
```

Prototype: https://github.com/ExpHP/numpy/commit/a3507d319dd39fc68bef0131df84103f102f2f15

(note: In this design, the partially-initialized decorator can only be used once.  I believe making it a pure function is possible, but this requires a bit more complexity, so I figured one-time use is a reasonable limitation for something intended for use in decorator syntax!)
@ExpHP OK, that looks pretty reasonable to me. Can you measure the difference in call-time overhead for a trivial function, e.g., calling your `identity` function in a loop? I think this will add a small amount of overhead, but hopefully not too much.

(It may seem a little excessive to insist on performance here, but `vectorize` already has some performance optimizations so clearly someone cared about this.)
> @ExpHP OK, that looks pretty reasonable to me. Can you measure the difference in call-time overhead for a trivial function, e.g., calling your identity function in a loop? I think this will add a small amount of overhead, but hopefully not too much.

I'll take a look.  If it turns out to be an issue, my initial design for the implementation looked something more like this, which effectively adds zero overhead to `__call__` once it is finally initialized, but is perhaps a bit "evil" in design:

```python
class vectorize:

    def __init__(self, pyfunc=None, ...):
        ...
        # current definition of __init__
        ...
        if pyfunc:
            self(pyfunc)

    def _call_as_decorator(self, pyfunc):
        self.pyfunc = pyfunc
        self.__call__ = self._call_as_normal
        return self

    def _call_as_normal(self, *args, **kw):
        ...
        # current definition of __call__
        ...

    __call__ = _call_as_decorator
```
Hmm.  Yes, looks like the additional overhead of the `if` in `__call__` costs **~0.4s per 1,000,000 calls.**

The most trivial case I can think of is `np.vectorize((lambda x: x), otypes=['object'])` called on a scalar, which currently takes **~9.4s** per 1,000,000 calls. This puts the maximum theoretical slowdown at about **5%**.

Benchmark script and output: https://gist.github.com/ExpHP/9b65c918e8e5a5fc2ae5dfe63f7675cf
I think I'd prefer overriding `vectorize.__new__`~ to return a lambda if the function argument is missing, otherwise fall back on `object.__new__`

Commit message: Added a test for positional args (PR-23061)

Patch to review:
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 7a69c3c81456..30349f9e5868 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -2119,10 +2119,10 @@ def _create_arrays(broadcast_shape, dim_sizes, list_of_core_dims, dtypes,
 @set_module('numpy')
 class vectorize:
     """
-    vectorize(pyfunc, otypes=None, doc=None, excluded=None, cache=False,
-              signature=None)
+    vectorize(pyfunc=np._NoValue, otypes=None, doc=None, excluded=None,
+    cache=False, signature=None)
 
-    Generalized function class.
+    Returns an object that acts like pyfunc, but takes arrays as input.
 
     Define a vectorized function which takes a nested sequence of objects or
     numpy arrays as inputs and returns a single numpy array or a tuple of numpy
@@ -2136,8 +2136,9 @@ class vectorize:
 
     Parameters
     ----------
-    pyfunc : callable
+    pyfunc : callable, optional
         A python function or method.
+        Can be omitted to produce a decorator with keyword arguments.
     otypes : str or list of dtypes, optional
         The output data type. It must be specified as either a string of
         typecode characters or a list of data type specifiers. There should
@@ -2169,8 +2170,9 @@ class vectorize:
 
     Returns
     -------
-    vectorized : callable
-        Vectorized function.
+    out : callable
+        A vectorized function if ``pyfunc`` was provided,
+        a decorator otherwise.
 
     See Also
     --------
@@ -2267,18 +2269,55 @@ class vectorize:
            [0., 0., 1., 2., 1., 0.],
            [0., 0., 0., 1., 2., 1.]])
 
+    Decorator syntax is supported.  The decorator can be called as
+    a function to provide keyword arguments.
+    >>>@np.vectorize
+    ...def identity(x):
+    ...    return x
+    ...
+    >>>identity([0, 1, 2])
+    array([0, 1, 2])
+    >>>@np.vectorize(otypes=[float])
+    ...def as_float(x):
+    ...    return x
+    ...
+    >>>as_float([0, 1, 2])
+    array([0., 1., 2.])
     """
-    def __init__(self, pyfunc, otypes=None, doc=None, excluded=None,
-                 cache=False, signature=None):
+    def __init__(self, pyfunc=np._NoValue, otypes=None, doc=None,
+                 excluded=None, cache=False, signature=None):
+
+        if not callable(pyfunc):
+            p_temp = pyfunc
+            pyfunc = np._NoValue
+            if p_temp is not None and p_temp is not np._NoValue:
+                o_temp = otypes
+                otypes = p_temp
+                if o_temp is not None:
+                    d_temp = doc
+                    doc = o_temp
+                    if d_temp is not None:
+                        e_temp = excluded
+                        excluded = d_temp
+                        if e_temp is True or e_temp is False:
+                            c_temp = cache
+                            cache = e_temp
+                            if c_temp is not None:
+                                signature = c_temp
+
         self.pyfunc = pyfunc
         self.cache = cache
         self.signature = signature
-        self._ufunc = {}    # Caching to improve default performance
+        if pyfunc != np._NoValue:
+            self.__name__ = pyfunc.__name__
 
+        self._ufunc = {}    # Caching to improve default performance
+        self._doc = None
+        self.__doc__ = doc
         if doc is None:
             self.__doc__ = pyfunc.__doc__
         else:
-            self.__doc__ = doc
+            self._doc = doc
 
         if isinstance(otypes, str):
             for char in otypes:
@@ -2300,7 +2339,15 @@ def __init__(self, pyfunc, otypes=None, doc=None, excluded=None,
         else:
             self._in_and_out_core_dims = None
 
-    def __call__(self, *args, **kwargs):
+    def _init_stage_2(self, pyfunc, *args, **kwargs):
+        self.__name__ = pyfunc.__name__
+        self.pyfunc = pyfunc
+        if self._doc is None:
+            self.__doc__ = pyfunc.__doc__
+        else:
+            self.__doc__ = self._doc
+
+    def _call_as_normal(self, *args, **kwargs):
         """
         Return arrays with the results of `pyfunc` broadcast (vectorized) over
         `args` and `kwargs` not in `excluded`.
@@ -2330,6 +2377,13 @@ def func(*vargs):
 
         return self._vectorize_call(func=func, args=vargs)
 
+    def __call__(self, *args, **kwargs):
+        if self.pyfunc is np._NoValue:
+            self._init_stage_2(*args, **kwargs)
+            return self
+
+        return self._call_as_normal(*args, **kwargs)
+
     def _get_ufunc_and_otypes(self, func, args):
         """Return (ufunc, otypes)."""
         # frompyfunc will fail if args is empty
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index ecba35f2d39f..b56d776cbc2c 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1780,6 +1780,62 @@ class subclass(np.ndarray):
         assert_equal(type(r), subclass)
         assert_equal(r, m * v)
 
+    def test_name(self):
+        #See gh-23021
+        @np.vectorize
+        def f2(a, b):
+            return a + b
+
+        assert f2.__name__ == 'f2'
+
+    def test_decorator(self):
+        @vectorize
+        def addsubtract(a, b):
+            if a > b:
+                return a - b
+            else:
+                return a + b
+
+        r = addsubtract([0, 3, 6, 9], [1, 3, 5, 7])
+        assert_array_equal(r, [1, 6, 1, 2])
+
+    def test_docstring(self):
+        @vectorize
+        def f(x):
+            """Docstring"""
+            return x
+
+        assert f.__doc__ == "Docstring"
+
+    def test_signature_otypes_decorator(self):
+        @vectorize(signature='(n)->(n)', otypes=['float64'])
+        def f(x):
+            return x
+
+        r = f([1, 2, 3])
+        assert_equal(r.dtype, np.dtype('float64'))
+        assert_array_equal(r, [1, 2, 3])
+        assert f.__name__ == 'f'
+
+    def test_decorator_positional_args(self):
+        A = np.vectorize(abs, ['float64'], "return float absolute value")
+        y1 = A([1, -1, 0, 0.3, 5])
+
+        @vectorize(['float64'], "return float absolute value")
+        def myabs(a):
+            return abs(a)
+
+        y2 = myabs([1, -1, 0, 0.3, 5])
+        assert_array_equal(y1, y2)
+
+    def test_positional_regression_9477(self):
+        # This supplies the first keyword argument as a positional,
+        # to ensure that they are still properly forwarded after the
+        # enhancement for #9477
+        f = vectorize((lambda x: x), ['float64'])
+        r = f([2])
+        assert_equal(r.dtype, np.dtype('float64'))
+
 
 class TestLeaks:
     class A:


Code review comments:
diff_hunk: @@ -2286,6 +2286,25 @@ class vectorize:
     """
     def __init__(self, pyfunc=np._NoValue, otypes=None, doc=None,
                  excluded=None, cache=False, signature=None):
+
+        if not callable(pyfunc):
+            p_temp = pyfunc
+            pyfunc = np._NoValue
+            if p_temp is not None and p_temp is not np._NoValue:
+                o_temp = otypes
+                otypes = p_temp
path: numpy/lib/function_base.py
text: @user1:
What if someone has given really wrong input, like `vectorize(1, otypes=[...])`. I fear that for the shifting, one needs to do a `*args, **kwargs` and re-call (or re-init) with `None, *args, **kwargs` so that python can remove duplications...

@author:
>What if someone has given really wrong input

That would be assuming that the user is trying to use vectorize without reading the documentation, which I don't think is right. Anyways, from tinkering with the code while trying to solve this, it seems that what would happen is that vectorize would raise ```ValueError("Invalid otype specification")```. There might be some specific incorrect input that avoids this, but I don't think that's something we have to be concerned about.