1"""
2Methods that can be shared by many array-like classes or subclasses:
3    Series
4    Index
5    ExtensionArray
6"""
7import operator
8from typing import Any, Callable
9import warnings
10
11import numpy as np
12
13from pandas._libs import lib
14
15from pandas.core.construction import extract_array
16from pandas.core.ops import maybe_dispatch_ufunc_to_dunder_op, roperator
17from pandas.core.ops.common import unpack_zerodim_and_defer
18
19
20class OpsMixin:
21    # -------------------------------------------------------------
22    # Comparisons
23
24    def _cmp_method(self, other, op):
25        return NotImplemented
26
27    @unpack_zerodim_and_defer("__eq__")
28    def __eq__(self, other):
29        return self._cmp_method(other, operator.eq)
30
31    @unpack_zerodim_and_defer("__ne__")
32    def __ne__(self, other):
33        return self._cmp_method(other, operator.ne)
34
35    @unpack_zerodim_and_defer("__lt__")
36    def __lt__(self, other):
37        return self._cmp_method(other, operator.lt)
38
39    @unpack_zerodim_and_defer("__le__")
40    def __le__(self, other):
41        return self._cmp_method(other, operator.le)
42
43    @unpack_zerodim_and_defer("__gt__")
44    def __gt__(self, other):
45        return self._cmp_method(other, operator.gt)
46
47    @unpack_zerodim_and_defer("__ge__")
48    def __ge__(self, other):
49        return self._cmp_method(other, operator.ge)
50
51    # -------------------------------------------------------------
52    # Logical Methods
53
54    def _logical_method(self, other, op):
55        return NotImplemented
56
57    @unpack_zerodim_and_defer("__and__")
58    def __and__(self, other):
59        return self._logical_method(other, operator.and_)
60
61    @unpack_zerodim_and_defer("__rand__")
62    def __rand__(self, other):
63        return self._logical_method(other, roperator.rand_)
64
65    @unpack_zerodim_and_defer("__or__")
66    def __or__(self, other):
67        return self._logical_method(other, operator.or_)
68
69    @unpack_zerodim_and_defer("__ror__")
70    def __ror__(self, other):
71        return self._logical_method(other, roperator.ror_)
72
73    @unpack_zerodim_and_defer("__xor__")
74    def __xor__(self, other):
75        return self._logical_method(other, operator.xor)
76
77    @unpack_zerodim_and_defer("__rxor__")
78    def __rxor__(self, other):
79        return self._logical_method(other, roperator.rxor)
80
81    # -------------------------------------------------------------
82    # Arithmetic Methods
83
84    def _arith_method(self, other, op):
85        return NotImplemented
86
87    @unpack_zerodim_and_defer("__add__")
88    def __add__(self, other):
89        return self._arith_method(other, operator.add)
90
91    @unpack_zerodim_and_defer("__radd__")
92    def __radd__(self, other):
93        return self._arith_method(other, roperator.radd)
94
95    @unpack_zerodim_and_defer("__sub__")
96    def __sub__(self, other):
97        return self._arith_method(other, operator.sub)
98
99    @unpack_zerodim_and_defer("__rsub__")
100    def __rsub__(self, other):
101        return self._arith_method(other, roperator.rsub)
102
103    @unpack_zerodim_and_defer("__mul__")
104    def __mul__(self, other):
105        return self._arith_method(other, operator.mul)
106
107    @unpack_zerodim_and_defer("__rmul__")
108    def __rmul__(self, other):
109        return self._arith_method(other, roperator.rmul)
110
111    @unpack_zerodim_and_defer("__truediv__")
112    def __truediv__(self, other):
113        return self._arith_method(other, operator.truediv)
114
115    @unpack_zerodim_and_defer("__rtruediv__")
116    def __rtruediv__(self, other):
117        return self._arith_method(other, roperator.rtruediv)
118
119    @unpack_zerodim_and_defer("__floordiv__")
120    def __floordiv__(self, other):
121        return self._arith_method(other, operator.floordiv)
122
123    @unpack_zerodim_and_defer("__rfloordiv")
124    def __rfloordiv__(self, other):
125        return self._arith_method(other, roperator.rfloordiv)
126
127    @unpack_zerodim_and_defer("__mod__")
128    def __mod__(self, other):
129        return self._arith_method(other, operator.mod)
130
131    @unpack_zerodim_and_defer("__rmod__")
132    def __rmod__(self, other):
133        return self._arith_method(other, roperator.rmod)
134
135    @unpack_zerodim_and_defer("__divmod__")
136    def __divmod__(self, other):
137        return self._arith_method(other, divmod)
138
139    @unpack_zerodim_and_defer("__rdivmod__")
140    def __rdivmod__(self, other):
141        return self._arith_method(other, roperator.rdivmod)
142
143    @unpack_zerodim_and_defer("__pow__")
144    def __pow__(self, other):
145        return self._arith_method(other, operator.pow)
146
147    @unpack_zerodim_and_defer("__rpow__")
148    def __rpow__(self, other):
149        return self._arith_method(other, roperator.rpow)
150
151
152# -----------------------------------------------------------------------------
153# Helpers to implement __array_ufunc__
154
155
156def _is_aligned(frame, other):
157    """
158    Helper to check if a DataFrame is aligned with another DataFrame or Series.
159    """
160    from pandas import DataFrame
161
162    if isinstance(other, DataFrame):
163        return frame._indexed_same(other)
164    else:
165        # Series -> match index
166        return frame.columns.equals(other.index)
167
168
169def _maybe_fallback(ufunc: Callable, method: str, *inputs: Any, **kwargs: Any):
170    """
171    In the future DataFrame, inputs to ufuncs will be aligned before applying
172    the ufunc, but for now we ignore the index but raise a warning if behaviour
173    would change in the future.
174    This helper detects the case where a warning is needed and then fallbacks
175    to applying the ufunc on arrays to avoid alignment.
176
177    See https://github.com/pandas-dev/pandas/pull/39239
178    """
179    from pandas import DataFrame
180    from pandas.core.generic import NDFrame
181
182    n_alignable = sum(isinstance(x, NDFrame) for x in inputs)
183    n_frames = sum(isinstance(x, DataFrame) for x in inputs)
184
185    if n_alignable >= 2 and n_frames >= 1:
186        # if there are 2 alignable inputs (Series or DataFrame), of which at least 1
187        # is a DataFrame -> we would have had no alignment before -> warn that this
188        # will align in the future
189
190        # the first frame is what determines the output index/columns in pandas < 1.2
191        first_frame = next(x for x in inputs if isinstance(x, DataFrame))
192
193        # check if the objects are aligned or not
194        non_aligned = sum(
195            not _is_aligned(first_frame, x) for x in inputs if isinstance(x, NDFrame)
196        )
197
198        # if at least one is not aligned -> warn and fallback to array behaviour
199        if non_aligned:
200            warnings.warn(
201                "Calling a ufunc on non-aligned DataFrames (or DataFrame/Series "
202                "combination). Currently, the indices are ignored and the result "
203                "takes the index/columns of the first DataFrame. In the future , "
204                "the DataFrames/Series will be aligned before applying the ufunc.\n"
205                "Convert one of the arguments to a NumPy array "
206                "(eg 'ufunc(df1, np.asarray(df2)') to keep the current behaviour, "
207                "or align manually (eg 'df1, df2 = df1.align(df2)') before passing to "
208                "the ufunc to obtain the future behaviour and silence this warning.",
209                FutureWarning,
210                stacklevel=4,
211            )
212
213            # keep the first dataframe of the inputs, other DataFrame/Series is
214            # converted to array for fallback behaviour
215            new_inputs = []
216            for x in inputs:
217                if x is first_frame:
218                    new_inputs.append(x)
219                elif isinstance(x, NDFrame):
220                    new_inputs.append(np.asarray(x))
221                else:
222                    new_inputs.append(x)
223
224            # call the ufunc on those transformed inputs
225            return getattr(ufunc, method)(*new_inputs, **kwargs)
226
227    # signal that we didn't fallback / execute the ufunc yet
228    return NotImplemented
229
230
231def array_ufunc(self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any):
232    """
233    Compatibility with numpy ufuncs.
234
235    See also
236    --------
237    numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__
238    """
239    from pandas.core.generic import NDFrame
240    from pandas.core.internals import BlockManager
241
242    cls = type(self)
243
244    # for backwards compatibility check and potentially fallback for non-aligned frames
245    result = _maybe_fallback(ufunc, method, *inputs, **kwargs)
246    if result is not NotImplemented:
247        return result
248
249    # for binary ops, use our custom dunder methods
250    result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs)
251    if result is not NotImplemented:
252        return result
253
254    # Determine if we should defer.
255    no_defer = (np.ndarray.__array_ufunc__, cls.__array_ufunc__)
256
257    for item in inputs:
258        higher_priority = (
259            hasattr(item, "__array_priority__")
260            and item.__array_priority__ > self.__array_priority__
261        )
262        has_array_ufunc = (
263            hasattr(item, "__array_ufunc__")
264            and type(item).__array_ufunc__ not in no_defer
265            and not isinstance(item, self._HANDLED_TYPES)
266        )
267        if higher_priority or has_array_ufunc:
268            return NotImplemented
269
270    # align all the inputs.
271    types = tuple(type(x) for x in inputs)
272    alignable = [x for x, t in zip(inputs, types) if issubclass(t, NDFrame)]
273
274    if len(alignable) > 1:
275        # This triggers alignment.
276        # At the moment, there aren't any ufuncs with more than two inputs
277        # so this ends up just being x1.index | x2.index, but we write
278        # it to handle *args.
279
280        if len(set(types)) > 1:
281            # We currently don't handle ufunc(DataFrame, Series)
282            # well. Previously this raised an internal ValueError. We might
283            # support it someday, so raise a NotImplementedError.
284            raise NotImplementedError(
285                "Cannot apply ufunc {} to mixed DataFrame and Series "
286                "inputs.".format(ufunc)
287            )
288        axes = self.axes
289        for obj in alignable[1:]:
290            # this relies on the fact that we aren't handling mixed
291            # series / frame ufuncs.
292            for i, (ax1, ax2) in enumerate(zip(axes, obj.axes)):
293                axes[i] = ax1.union(ax2)
294
295        reconstruct_axes = dict(zip(self._AXIS_ORDERS, axes))
296        inputs = tuple(
297            x.reindex(**reconstruct_axes) if issubclass(t, NDFrame) else x
298            for x, t in zip(inputs, types)
299        )
300    else:
301        reconstruct_axes = dict(zip(self._AXIS_ORDERS, self.axes))
302
303    if self.ndim == 1:
304        names = [getattr(x, "name") for x in inputs if hasattr(x, "name")]
305        name = names[0] if len(set(names)) == 1 else None
306        reconstruct_kwargs = {"name": name}
307    else:
308        reconstruct_kwargs = {}
309
310    def reconstruct(result):
311        if lib.is_scalar(result):
312            return result
313        if result.ndim != self.ndim:
314            if method == "outer":
315                if self.ndim == 2:
316                    # we already deprecated for Series
317                    msg = (
318                        "outer method for ufunc {} is not implemented on "
319                        "pandas objects. Returning an ndarray, but in the "
320                        "future this will raise a 'NotImplementedError'. "
321                        "Consider explicitly converting the DataFrame "
322                        "to an array with '.to_numpy()' first."
323                    )
324                    warnings.warn(msg.format(ufunc), FutureWarning, stacklevel=4)
325                    return result
326                raise NotImplementedError
327            return result
328        if isinstance(result, BlockManager):
329            # we went through BlockManager.apply
330            result = self._constructor(result, **reconstruct_kwargs, copy=False)
331        else:
332            # we converted an array, lost our axes
333            result = self._constructor(
334                result, **reconstruct_axes, **reconstruct_kwargs, copy=False
335            )
336        # TODO: When we support multiple values in __finalize__, this
337        # should pass alignable to `__fianlize__` instead of self.
338        # Then `np.add(a, b)` would consider attrs from both a and b
339        # when a and b are NDFrames.
340        if len(alignable) == 1:
341            result = result.__finalize__(self)
342        return result
343
344    if self.ndim > 1 and (
345        len(inputs) > 1 or ufunc.nout > 1  # type: ignore[attr-defined]
346    ):
347        # Just give up on preserving types in the complex case.
348        # In theory we could preserve them for them.
349        # * nout>1 is doable if BlockManager.apply took nout and
350        #   returned a Tuple[BlockManager].
351        # * len(inputs) > 1 is doable when we know that we have
352        #   aligned blocks / dtypes.
353        inputs = tuple(np.asarray(x) for x in inputs)
354        result = getattr(ufunc, method)(*inputs, **kwargs)
355    elif self.ndim == 1:
356        # ufunc(series, ...)
357        inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs)
358        result = getattr(ufunc, method)(*inputs, **kwargs)
359    else:
360        # ufunc(dataframe)
361        if method == "__call__" and not kwargs:
362            # for np.<ufunc>(..) calls
363            # kwargs cannot necessarily be handled block-by-block, so only
364            # take this path if there are no kwargs
365            mgr = inputs[0]._mgr
366            result = mgr.apply(getattr(ufunc, method))
367        else:
368            # otherwise specific ufunc methods (eg np.<ufunc>.accumulate(..))
369            # Those can have an axis keyword and thus can't be called block-by-block
370            result = getattr(ufunc, method)(np.asarray(inputs[0]), **kwargs)
371
372    if ufunc.nout > 1:  # type: ignore[attr-defined]
373        result = tuple(reconstruct(x) for x in result)
374    else:
375        result = reconstruct(result)
376    return result
377