Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/arraylike.py: 35%

1""" 

2Methods that can be shared by many array-like classes or subclasses: 

3 Series 

4 Index 

5 ExtensionArray 

6""" 

7 

8from __future__ import annotations 

9 

10import operator 

11from typing import ( 

12 Any, 

13 Self, 

14) 

15 

16import numpy as np 

17 

18from pandas._libs import lib 

19from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op 

20 

21from pandas.core.dtypes.generic import ABCNDFrame 

22 

23from pandas.core import roperator 

24from pandas.core.construction import extract_array 

25from pandas.core.ops.common import unpack_zerodim_and_defer 

26from pandas.core.ops.invalid import disallow_datetimelike_logical_ufunc 

27 

28REDUCTION_ALIASES = { 

29 "maximum": "max", 

30 "minimum": "min", 

31 "add": "sum", 

32 "multiply": "prod", 

33} 

34 

35 

36class OpsMixin: 

37 # ------------------------------------------------------------- 

38 # Comparisons 

39 

40 def _cmp_method(self, other, op): 

41 return NotImplemented 

42 

43 @unpack_zerodim_and_defer("__eq__") 

44 def __eq__(self, other) -> Self: # type: ignore[override] 

45 return self._cmp_method(other, operator.eq) 

46 

47 @unpack_zerodim_and_defer("__ne__") 

48 def __ne__(self, other) -> Self: # type: ignore[override] 

49 return self._cmp_method(other, operator.ne) 

50 

51 @unpack_zerodim_and_defer("__lt__") 

52 def __lt__(self, other) -> Self: 

53 return self._cmp_method(other, operator.lt) 

54 

55 @unpack_zerodim_and_defer("__le__") 

56 def __le__(self, other) -> Self: 

57 return self._cmp_method(other, operator.le) 

58 

59 @unpack_zerodim_and_defer("__gt__") 

60 def __gt__(self, other) -> Self: 

61 return self._cmp_method(other, operator.gt) 

62 

63 @unpack_zerodim_and_defer("__ge__") 

64 def __ge__(self, other) -> Self: 

65 return self._cmp_method(other, operator.ge) 

66 

67 # ------------------------------------------------------------- 

68 # Logical Methods 

69 

70 def _logical_method(self, other, op): 

71 return NotImplemented 

72 

73 @unpack_zerodim_and_defer("__and__") 

74 def __and__(self, other) -> Self: 

75 """ 

76 Return elementwise ``self & other``. 

77 

78 Logical AND for boolean operands, bitwise AND for integer operands. 

79 """ 

80 return self._logical_method(other, operator.and_) 

81 

82 @unpack_zerodim_and_defer("__rand__") 

83 def __rand__(self, other) -> Self: 

84 """Return elementwise ``other & self``.""" 

85 return self._logical_method(other, roperator.rand_) 

86 

87 @unpack_zerodim_and_defer("__or__") 

88 def __or__(self, other) -> Self: 

89 """ 

90 Return elementwise ``self | other``. 

91 

92 Logical OR for boolean operands, bitwise OR for integer operands. 

93 """ 

94 return self._logical_method(other, operator.or_) 

95 

96 @unpack_zerodim_and_defer("__ror__") 

97 def __ror__(self, other) -> Self: 

98 """Return elementwise ``other | self``.""" 

99 return self._logical_method(other, roperator.ror_) 

100 

101 @unpack_zerodim_and_defer("__xor__") 

102 def __xor__(self, other) -> Self: 

103 """ 

104 Return elementwise ``self ^ other``. 

105 

106 Logical XOR for boolean operands, bitwise XOR for integer operands. 

107 """ 

108 return self._logical_method(other, operator.xor) 

109 

110 @unpack_zerodim_and_defer("__rxor__") 

111 def __rxor__(self, other) -> Self: 

112 """Return elementwise ``other ^ self``.""" 

113 return self._logical_method(other, roperator.rxor) 

114 

115 # ------------------------------------------------------------- 

116 # Arithmetic Methods 

117 

118 def _arith_method(self, other, op): 

119 return NotImplemented 

120 

121 @unpack_zerodim_and_defer("__add__") 

122 def __add__(self, other) -> Self: 

123 """ 

124 Get Addition of DataFrame and other, column-wise. 

125 

126 Equivalent to ``DataFrame.add(other)``. 

127 

128 Parameters 

129 ---------- 

130 other : scalar, sequence, Series, dict or DataFrame 

131 Object to be added to the DataFrame. 

132 

133 Returns 

134 ------- 

135 DataFrame 

136 The result of adding ``other`` to DataFrame. 

137 

138 See Also 

139 -------- 

140 DataFrame.add : Add a DataFrame and another object, with option for index- 

141 or column-oriented addition. 

142 

143 Examples 

144 -------- 

145 >>> df = pd.DataFrame( 

146 ... {"height": [1.5, 2.6], "weight": [500, 800]}, index=["elk", "moose"] 

147 ... ) 

148 >>> df 

149 height weight 

150 elk 1.5 500 

151 moose 2.6 800 

152 

153 Adding a scalar affects all rows and columns. 

154 

155 >>> df[["height", "weight"]] + 1.5 

156 height weight 

157 elk 3.0 501.5 

158 moose 4.1 801.5 

159 

160 Each element of a list is added to a column of the DataFrame, in order. 

161 

162 >>> df[["height", "weight"]] + [0.5, 1.5] 

163 height weight 

164 elk 2.0 501.5 

165 moose 3.1 801.5 

166 

167 Keys of a dictionary are aligned to the DataFrame, based on column names; 

168 each value in the dictionary is added to the corresponding column. 

169 

170 >>> df[["height", "weight"]] + {"height": 0.5, "weight": 1.5} 

171 height weight 

172 elk 2.0 501.5 

173 moose 3.1 801.5 

174 

175 When `other` is a :class:`Series`, the index of `other` is aligned with the 

176 columns of the DataFrame. 

177 

178 >>> s1 = pd.Series([0.5, 1.5], index=["weight", "height"]) 

179 >>> df[["height", "weight"]] + s1 

180 height weight 

181 elk 3.0 500.5 

182 moose 4.1 800.5 

183 

184 Even when the index of `other` is the same as the index of the DataFrame, 

185 the :class:`Series` will not be reoriented. If index-wise alignment is desired, 

186 :meth:`DataFrame.add` should be used with `axis='index'`. 

187 

188 >>> s2 = pd.Series([0.5, 1.5], index=["elk", "moose"]) 

189 >>> df[["height", "weight"]] + s2 

190 elk height moose weight 

191 elk NaN NaN NaN NaN 

192 moose NaN NaN NaN NaN 

193 

194 >>> df[["height", "weight"]].add(s2, axis="index") 

195 height weight 

196 elk 2.0 500.5 

197 moose 4.1 801.5 

198 

199 When `other` is a :class:`DataFrame`, both columns names and the 

200 index are aligned. 

201 

202 >>> other = pd.DataFrame( 

203 ... {"height": [0.2, 0.4, 0.6]}, index=["elk", "moose", "deer"] 

204 ... ) 

205 >>> df[["height", "weight"]] + other 

206 height weight 

207 deer NaN NaN 

208 elk 1.7 NaN 

209 moose 3.0 NaN 

210 """ 

211 return self._arith_method(other, operator.add) 

212 

213 @unpack_zerodim_and_defer("__radd__") 

214 def __radd__(self, other) -> Self: 

215 return self._arith_method(other, roperator.radd) 

216 

217 @unpack_zerodim_and_defer("__sub__") 

218 def __sub__(self, other) -> Self: 

219 return self._arith_method(other, operator.sub) 

220 

221 @unpack_zerodim_and_defer("__rsub__") 

222 def __rsub__(self, other) -> Self: 

223 return self._arith_method(other, roperator.rsub) 

224 

225 @unpack_zerodim_and_defer("__mul__") 

226 def __mul__(self, other) -> Self: 

227 return self._arith_method(other, operator.mul) 

228 

229 @unpack_zerodim_and_defer("__rmul__") 

230 def __rmul__(self, other) -> Self: 

231 return self._arith_method(other, roperator.rmul) 

232 

233 @unpack_zerodim_and_defer("__truediv__") 

234 def __truediv__(self, other) -> Self: 

235 return self._arith_method(other, operator.truediv) 

236 

237 @unpack_zerodim_and_defer("__rtruediv__") 

238 def __rtruediv__(self, other) -> Self: 

239 return self._arith_method(other, roperator.rtruediv) 

240 

241 @unpack_zerodim_and_defer("__floordiv__") 

242 def __floordiv__(self, other) -> Self: 

243 return self._arith_method(other, operator.floordiv) 

244 

245 @unpack_zerodim_and_defer("__rfloordiv") 

246 def __rfloordiv__(self, other) -> Self: 

247 return self._arith_method(other, roperator.rfloordiv) 

248 

249 @unpack_zerodim_and_defer("__mod__") 

250 def __mod__(self, other) -> Self: 

251 return self._arith_method(other, operator.mod) 

252 

253 @unpack_zerodim_and_defer("__rmod__") 

254 def __rmod__(self, other) -> Self: 

255 return self._arith_method(other, roperator.rmod) 

256 

257 @unpack_zerodim_and_defer("__divmod__") 

258 def __divmod__(self, other) -> tuple[Self, Self]: 

259 return self._arith_method(other, divmod) 

260 

261 @unpack_zerodim_and_defer("__rdivmod__") 

262 def __rdivmod__(self, other) -> tuple[Self, Self]: 

263 return self._arith_method(other, roperator.rdivmod) 

264 

265 @unpack_zerodim_and_defer("__pow__") 

266 def __pow__(self, other) -> Self: 

267 return self._arith_method(other, operator.pow) 

268 

269 @unpack_zerodim_and_defer("__rpow__") 

270 def __rpow__(self, other) -> Self: 

271 return self._arith_method(other, roperator.rpow) 

272 

273 

274# ----------------------------------------------------------------------------- 

275# Helpers to implement __array_ufunc__ 

276 

277 

278def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any): 

279 """ 

280 Compatibility with numpy ufuncs. 

281 

282 See also 

283 -------- 

284 numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__ 

285 """ 

286 from pandas.core.frame import ( 

287 DataFrame, 

288 Series, 

289 ) 

290 from pandas.core.generic import NDFrame 

291 from pandas.core.internals import BlockManager 

292 

293 cls = type(self) 

294 

295 disallow_datetimelike_logical_ufunc(ufunc, inputs) 

296 

297 kwargs = _standardize_out_kwarg(**kwargs) 

298 

299 # for binary ops, use our custom dunder methods 

300 result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs) 

301 if result is not NotImplemented: 

302 return result 

303 

304 # Determine if we should defer. 

305 no_defer = ( 

306 np.ndarray.__array_ufunc__, 

307 cls.__array_ufunc__, 

308 ) 

309 

310 for item in inputs: 

311 higher_priority = ( 

312 hasattr(item, "__array_priority__") 

313 and item.__array_priority__ > self.__array_priority__ 

314 ) 

315 has_array_ufunc = ( 

316 hasattr(item, "__array_ufunc__") 

317 and type(item).__array_ufunc__ not in no_defer 

318 and not isinstance(item, self._HANDLED_TYPES) 

319 ) 

320 if higher_priority or has_array_ufunc: 

321 return NotImplemented 

322 

323 # align all the inputs. 

324 types = tuple(type(x) for x in inputs) 

325 alignable = [ 

326 x for x, t in zip(inputs, types, strict=True) if issubclass(t, NDFrame) 

327 ] 

328 

329 if len(alignable) > 1: 

330 # This triggers alignment. 

331 # At the moment, there aren't any ufuncs with more than two inputs 

332 # so this ends up just being x1.index | x2.index, but we write 

333 # it to handle *args. 

334 set_types = set(types) 

335 if len(set_types) > 1 and {DataFrame, Series}.issubset(set_types): 

336 # We currently don't handle ufunc(DataFrame, Series) 

337 # well. Previously this raised an internal ValueError. We might 

338 # support it someday, so raise a NotImplementedError. 

339 raise NotImplementedError( 

340 f"Cannot apply ufunc {ufunc} to mixed DataFrame and Series inputs." 

341 ) 

342 if len(alignable) == 2: 

343 # use align rather than union+reindex to support duplicate 

344 # labels, matching the dunder ops, see GH#54416 

345 left, right = alignable[0].align(alignable[1], join="outer") 

346 aligned = iter([left, right]) 

347 inputs = tuple( 

348 next(aligned) if issubclass(t, NDFrame) else x 

349 for x, t in zip(inputs, types, strict=True) 

350 ) 

351 axes = left.axes 

352 else: 

353 axes = self.axes 

354 for obj in alignable[1:]: 

355 # this relies on the fact that we aren't handling mixed 

356 # series / frame ufuncs. 

357 for i, (ax1, ax2) in enumerate(zip(axes, obj.axes, strict=True)): 

358 axes[i] = ax1.union(ax2) 

359 inputs = tuple( 

360 x.reindex(**dict(zip(self._AXIS_ORDERS, axes, strict=True))) 

361 if issubclass(t, NDFrame) 

362 else x 

363 for x, t in zip(inputs, types, strict=True) 

364 ) 

365 

366 reconstruct_axes = dict(zip(self._AXIS_ORDERS, axes, strict=True)) 

367 else: 

368 reconstruct_axes = dict(zip(self._AXIS_ORDERS, self.axes, strict=True)) 

369 

370 if self.ndim == 1: 

371 names = {x.name for x in inputs if hasattr(x, "name")} 

372 name = names.pop() if len(names) == 1 else None 

373 reconstruct_kwargs = {"name": name} 

374 else: 

375 reconstruct_kwargs = {} 

376 

377 def reconstruct(result): 

378 if ufunc.nout > 1: 

379 # np.modf, np.frexp, np.divmod 

380 return tuple(_reconstruct(x) for x in result) 

381 

382 return _reconstruct(result) 

383 

384 def _reconstruct(result): 

385 if lib.is_scalar(result): 

386 return result 

387 

388 if result.ndim != self.ndim: 

389 if method == "outer": 

390 raise NotImplementedError 

391 return result 

392 if isinstance(result, BlockManager): 

393 # we went through BlockManager.apply e.g. np.sqrt 

394 result = self._constructor_from_mgr(result, axes=result.axes) 

395 else: 

396 # we converted an array, lost our axes 

397 result = self._constructor( 

398 result, **reconstruct_axes, **reconstruct_kwargs, copy=False 

399 ) 

400 # TODO: When we support multiple values in __finalize__, this 

401 # should pass alignable to `__finalize__` instead of self. 

402 # Then `np.add(a, b)` would consider attrs from both a and b 

403 # when a and b are NDFrames. 

404 if len(alignable) == 1: 

405 result = result.__finalize__(self) 

406 return result 

407 

408 if "out" in kwargs: 

409 # e.g. test_multiindex_get_loc 

410 result = dispatch_ufunc_with_out(self, ufunc, method, *inputs, **kwargs) 

411 return reconstruct(result) 

412 

413 if method == "reduce": 

414 # e.g. test.series.test_ufunc.test_reduce 

415 result = dispatch_reduction_ufunc(self, ufunc, method, *inputs, **kwargs) 

416 if result is not NotImplemented: 

417 return result 

418 

419 # We still get here with kwargs `axis` for e.g. np.maximum.accumulate 

420 # and `dtype` and `keepdims` for np.ptp 

421 

422 if self.ndim > 1 and (len(inputs) > 1 or ufunc.nout > 1): 

423 # Just give up on preserving types in the complex case. 

424 # In theory we could preserve them for them. 

425 # * nout>1 is doable if BlockManager.apply took nout and 

426 # returned a Tuple[BlockManager]. 

427 # * len(inputs) > 1 is doable when we know that we have 

428 # aligned blocks / dtypes. 

429 

430 # e.g. my_ufunc, modf, logaddexp, heaviside, subtract, add 

431 inputs = tuple(np.asarray(x) for x in inputs) 

432 # Note: we can't use default_array_ufunc here bc reindexing means 

433 # that `self` may not be among `inputs` 

434 result = getattr(ufunc, method)(*inputs, **kwargs) 

435 elif self.ndim == 1: 

436 # ufunc(series, ...) 

437 inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs) 

438 result = getattr(ufunc, method)(*inputs, **kwargs) 

439 # ufunc(dataframe) 

440 elif method == "__call__" and not kwargs: 

441 # for np.<ufunc>(..) calls 

442 # kwargs cannot necessarily be handled block-by-block, so only 

443 # take this path if there are no kwargs 

444 mgr = inputs[0]._mgr # pyright: ignore[reportGeneralTypeIssues] 

445 result = mgr.apply(getattr(ufunc, method)) 

446 else: 

447 # otherwise specific ufunc methods (eg np.<ufunc>.accumulate(..)) 

448 # Those can have an axis keyword and thus can't be called block-by-block 

449 result = default_array_ufunc(inputs[0], ufunc, method, *inputs, **kwargs) # pyright: ignore[reportGeneralTypeIssues] 

450 # e.g. np.negative (only one reached), with "where" and "out" in kwargs 

451 

452 result = reconstruct(result) 

453 return result 

454 

455 

456def _standardize_out_kwarg(**kwargs) -> dict[str, Any]: 

457 """ 

458 If kwargs contain "out1" and "out2", replace that with a tuple "out" 

459 

460 np.divmod, np.modf, np.frexp can have either `out=(out1, out2)` or 

461 `out1=out1, out2=out2)` 

462 """ 

463 if "out" not in kwargs and "out1" in kwargs and "out2" in kwargs: 

464 out1 = kwargs.pop("out1") 

465 out2 = kwargs.pop("out2") 

466 out = (out1, out2) 

467 kwargs["out"] = out 

468 return kwargs 

469 

470 

471def dispatch_ufunc_with_out(self, ufunc: np.ufunc, method: str, *inputs, **kwargs): 

472 """ 

473 If we have an `out` keyword, then call the ufunc without `out` and then 

474 set the result into the given `out`. 

475 """ 

476 

477 # Note: we assume _standardize_out_kwarg has already been called. 

478 out = kwargs.pop("out") 

479 where = kwargs.pop("where", None) 

480 

481 result = getattr(ufunc, method)(*inputs, **kwargs) 

482 

483 if result is NotImplemented: 

484 return NotImplemented 

485 

486 if isinstance(result, tuple): 

487 # i.e. np.divmod, np.modf, np.frexp 

488 if not isinstance(out, tuple) or len(out) != len(result): 

489 raise NotImplementedError 

490 

491 for arr, res in zip(out, result, strict=True): 

492 _assign_where(arr, res, where) 

493 

494 return out 

495 

496 if isinstance(out, tuple): 

497 if len(out) == 1: 

498 out = out[0] 

499 else: 

500 raise NotImplementedError 

501 

502 _assign_where(out, result, where) 

503 return out 

504 

505 

506def _assign_where(out, result, where) -> None: 

507 """ 

508 Set a ufunc result into 'out', masking with a 'where' argument if necessary. 

509 """ 

510 if where is None: 

511 # no 'where' arg passed to ufunc 

512 out[:] = result 

513 else: 

514 np.putmask(out, where, result) 

515 

516 

517def default_array_ufunc(self, ufunc: np.ufunc, method: str, *inputs, **kwargs): 

518 """ 

519 Fallback to the behavior we would get if we did not define __array_ufunc__. 

520 

521 Notes 

522 ----- 

523 We are assuming that `self` is among `inputs`. 

524 """ 

525 if not any(x is self for x in inputs): 

526 raise NotImplementedError 

527 

528 new_inputs = [x if x is not self else np.asarray(x) for x in inputs] 

529 

530 return getattr(ufunc, method)(*new_inputs, **kwargs) 

531 

532 

533def dispatch_reduction_ufunc(self, ufunc: np.ufunc, method: str, *inputs, **kwargs): 

534 """ 

535 Dispatch ufunc reductions to self's reduction methods. 

536 """ 

537 assert method == "reduce" 

538 

539 if len(inputs) != 1 or inputs[0] is not self: 

540 return NotImplemented 

541 

542 if ufunc.__name__ not in REDUCTION_ALIASES: 

543 return NotImplemented 

544 

545 method_name = REDUCTION_ALIASES[ufunc.__name__] 

546 

547 # NB: we are assuming that min/max represent minimum/maximum methods, 

548 # which would not be accurate for e.g. Timestamp.min 

549 if not hasattr(self, method_name): 

550 return NotImplemented 

551 

552 if self.ndim > 1: 

553 if isinstance(self, ABCNDFrame): 

554 # TODO: test cases where this doesn't hold, i.e. 2D DTA/TDA 

555 kwargs["numeric_only"] = False 

556 

557 if "axis" not in kwargs: 

558 # For DataFrame reductions we don't want the default axis=0 

559 # Note: np.min is not a ufunc, but uses array_function_dispatch, 

560 # so calls DataFrame.min (without ever getting here) with the np.min 

561 # default of axis=None, which DataFrame.min catches and changes to axis=0. 

562 # np.minimum.reduce(df) gets here bc axis is not in kwargs, 

563 # so we set axis=0 to match the behavior of np.minimum.reduce(df.values) 

564 kwargs["axis"] = 0 

565 

566 # By default, numpy's reductions do not skip NaNs, so we have to 

567 # pass skipna=False 

568 return getattr(self, method_name)(skipna=False, **kwargs)