|
12 | 12 | import numpy as np
|
13 | 13 |
|
14 | 14 | from pandas._libs import (
|
| 15 | + iNaT, |
15 | 16 | lib,
|
16 | 17 | missing as libmissing,
|
17 | 18 | )
|
|
39 | 40 | is_bool,
|
40 | 41 | is_bool_dtype,
|
41 | 42 | is_dtype_equal,
|
| 43 | + is_float, |
42 | 44 | is_float_dtype,
|
43 | 45 | is_integer_dtype,
|
44 | 46 | is_list_like,
|
| 47 | + is_numeric_dtype, |
45 | 48 | is_object_dtype,
|
46 | 49 | is_scalar,
|
47 | 50 | is_string_dtype,
|
@@ -543,6 +546,48 @@ def _cmp_method(self, other, op) -> BooleanArray:
|
543 | 546 |
|
544 | 547 | return BooleanArray(result, mask, copy=False)
|
545 | 548 |
|
| 549 | + def _maybe_mask_result(self, result, mask, other, op_name: str): |
| 550 | + """ |
| 551 | + Parameters |
| 552 | + ---------- |
| 553 | + result : array-like |
| 554 | + mask : array-like bool |
| 555 | + other : scalar or array-like |
| 556 | + op_name : str |
| 557 | + """ |
| 558 | + # if we have a float operand we are by-definition |
| 559 | + # a float result |
| 560 | + # or our op is a divide |
| 561 | + if ( |
| 562 | + (is_float_dtype(other) or is_float(other)) |
| 563 | + or (op_name in ["rtruediv", "truediv"]) |
| 564 | + or (is_float_dtype(self.dtype) and is_numeric_dtype(result.dtype)) |
| 565 | + ): |
| 566 | + from pandas.core.arrays import FloatingArray |
| 567 | + |
| 568 | + return FloatingArray(result, mask, copy=False) |
| 569 | + |
| 570 | + elif is_bool_dtype(result): |
| 571 | + from pandas.core.arrays import BooleanArray |
| 572 | + |
| 573 | + return BooleanArray(result, mask, copy=False) |
| 574 | + |
| 575 | + elif result.dtype == "timedelta64[ns]": |
| 576 | + # e.g. test_numeric_arr_mul_tdscalar_numexpr_path |
| 577 | + from pandas.core.arrays import TimedeltaArray |
| 578 | + |
| 579 | + result[mask] = iNaT |
| 580 | + return TimedeltaArray._simple_new(result) |
| 581 | + |
| 582 | + elif is_integer_dtype(result): |
| 583 | + from pandas.core.arrays import IntegerArray |
| 584 | + |
| 585 | + return IntegerArray(result, mask, copy=False) |
| 586 | + |
| 587 | + else: |
| 588 | + result[mask] = np.nan |
| 589 | + return result |
| 590 | + |
546 | 591 | def isna(self) -> np.ndarray:
|
547 | 592 | return self._mask.copy()
|
548 | 593 |
|
|
0 commit comments