diff options
Diffstat (limited to 'numpy/lib/function_base.py')
| -rw-r--r-- | numpy/lib/function_base.py | 58 |
1 files changed, 46 insertions, 12 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 27ed79044..cba04765a 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -2119,8 +2119,8 @@ 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. @@ -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,19 +2269,36 @@ 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): self.pyfunc = pyfunc self.cache = cache self.signature = signature - self.__name__ = pyfunc.__name__ - 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: @@ -2301,7 +2320,15 @@ class vectorize: 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`. @@ -2331,6 +2358,13 @@ class vectorize: 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 @@ -2457,7 +2491,7 @@ class vectorize: if outputs is None: for result, core_dims in zip(results, output_core_dims): - _update_dim_sizes(dim_sizes, result, core_dims) + _update_dim_sizes(dim_sizes, result, core_dims) outputs = _create_arrays(broadcast_shape, dim_sizes, output_core_dims, otypes, results) |