def _maybe_mask_result(self, result, mask, other, op_name: str): """ Parameters ---------- result : array-like mask : array-like bool other : scalar or array-like op_name : str """ # if we have a float operand we are by-definition # a float result # or our op is a divide if (is_float_dtype(other) or is_float(other)) or ( op_name in ["rtruediv", "truediv"] ): from pandas.core.arrays import FloatingArray return FloatingArray(result, mask, copy=False) elif is_bool_dtype(result): return BooleanArray(result, mask, copy=False) elif is_integer_dtype(result): from pandas.core.arrays import IntegerArray return IntegerArray(result, mask, copy=False) else: result[mask] = np.nan return result
def test_div(left_array, right_array): result = left_array / right_array expected = FloatingArray( np.array( [1.0, np.inf, np.nan, 0.0, np.nan, np.nan, np.nan, np.nan, np.nan], dtype="float64", ), np.array([False, False, True, False, False, True, True, True, True]), ) tm.assert_extension_array_equal(result, expected)
@pytest.mark.parametrize( "data, dtype, expected", [ # Basic NumPy defaults. ([1, 2], None, IntegerArray._from_sequence([1, 2])), ([1, 2], object, PandasArray(np.array([1, 2], dtype=object))), ( [1, 2], np.dtype("float32"), PandasArray(np.array([1.0, 2.0], dtype=np.dtype("float32"))), ), (np.array([1, 2], dtype="int64"), None, IntegerArray._from_sequence([1, 2])), ( np.array([1.0, 2.0], dtype="float64"), None, FloatingArray._from_sequence([1.0, 2.0]), ), # String alias passes through to NumPy ([1, 2], "float32", PandasArray(np.array([1, 2], dtype="float32"))), # Period alias ( [pd.Period("2000", "D"), pd.Period("2001", "D")], "Period[D]", period_array(["2000", "2001"], freq="D"), ), # Period dtype ( [pd.Period("2000", "D")], pd.PeriodDtype("D"), period_array(["2000"], freq="D"), ),