From f9f28bafc4cda3a233a94a7209a8ddb7d7291e95 Mon Sep 17 00:00:00 2001 From: Simon Hawkins Date: Sat, 18 Jul 2020 11:36:58 +0100 Subject: [PATCH] REGR: np.argwhere on pd.Series raises ValueError --- pandas/core/generic.py | 42 +++++++++++++++++++++++++---- pandas/core/indexes/base.py | 2 +- pandas/core/indexes/datetimelike.py | 2 +- pandas/core/indexes/period.py | 9 ++++--- 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index e46fde1f59f16..d147052f0b978 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -1777,16 +1777,48 @@ def empty(self) -> bool_t: def __array__(self, dtype=None) -> np.ndarray: return np.asarray(self._values, dtype=dtype) - def __array_wrap__(self, result, context=None): + def __array_wrap__( + self, + result: np.ndarray, + context: Optional[Tuple[Callable, Tuple[Any, ...], int]] = None, + ): + """ + Gets called after a ufunc and other functions. + + Parameters + ---------- + result: np.ndarray + The result of the ufunc or other function called on the NumPy array + returned by __array__ + context: tuple of (func, tuple, int) + This parameter is returned by ufuncs as a 3-element tuple: (name of the + ufunc, arguments of the ufunc, domain of the ufunc), but is not set by + other numpy functions.q + + Notes + ----- + Series implements __array_ufunc_ so this not called for ufunc on Series. + """ result = lib.item_from_zerodim(result) if is_scalar(result): # e.g. we get here with np.ptp(series) # ptp also requires the item_from_zerodim return result - d = self._construct_axes_dict(self._AXIS_ORDERS, copy=False) - return self._constructor(result, **d).__finalize__( - self, method="__array_wrap__" - ) + + # if shape of result is the same as self we use the original axes + if result.shape == self.shape: + d = self._construct_axes_dict(self._AXIS_ORDERS, copy=False) + else: + d = dict() + + # if ndim of result is the same as self we return the same object + # otherwise we just return the NumPy array + if result.ndim == self.ndim: + return self._constructor(result, **d).__finalize__( + self, method="__array_wrap__" + ) + else: + return result # ideally we would define this to avoid the getattr checks, but # is slower diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 3dbee7d0929cb..3b0ed35ccea25 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -574,7 +574,7 @@ def __array__(self, dtype=None) -> np.ndarray: def __array_wrap__(self, result, context=None): """ - Gets called after a ufunc. + Gets called after a ufunc and other functions. """ result = lib.item_from_zerodim(result) if is_bool_dtype(result) or lib.is_scalar(result) or np.ndim(result) > 1: diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 15a7e25238983..b9515df004e99 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -112,7 +112,7 @@ def values(self): def __array_wrap__(self, result, context=None): """ - Gets called after a ufunc. + Gets called after a ufunc and other functions. """ result = lib.item_from_zerodim(result) if is_bool_dtype(result) or lib.is_scalar(result): diff --git a/pandas/core/indexes/period.py b/pandas/core/indexes/period.py index 03e11b652477f..1de738a8d6ab7 100644 --- a/pandas/core/indexes/period.py +++ b/pandas/core/indexes/period.py @@ -345,10 +345,13 @@ def _int64index(self) -> Int64Index: def __array_wrap__(self, result, context=None): """ - Gets called after a ufunc. Needs additional handling as - PeriodIndex stores internal data as int dtype + Gets called after a ufunc and other functions. - Replace this to __numpy_ufunc__ in future version + Needs additional handling as PeriodIndex stores internal data as int + dtype + + Replace this to __numpy_ufunc__ in future version and implement + __array_function__ for Indexes """ if isinstance(context, tuple) and len(context) > 0: func = context[0]