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