def test_numpy_array_equivalence(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
numpy_array = np.array([scalar] * len(data), dtype=data.dtype.numpy_dtype)
pd_array = pd.array(numpy_array, dtype=data.dtype)
result = op(data, numpy_array)
expected = op(data, pd_array)
if isinstance(expected, ExtensionArray):
tm.assert_extension_array_equal(result, expected)
else:
# TODO div still gives float ndarray -> remove this once we have Float EA
tm.assert_numpy_array_equal(result, expected)
def test_array_NA(data, all_arithmetic_operators):
if "truediv" in all_arithmetic_operators:
pytest.skip("division with pd.NA raises")
if "floordiv" in all_arithmetic_operators and is_numpy_dev:
pytest.skip("NumpyDev behavior GH#40874")
data, _ = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
scalar = pd.NA
scalar_array = pd.array([pd.NA] * len(data), dtype=data.dtype)
result = op(data, scalar)
expected = op(data, scalar_array)
tm.assert_extension_array_equal(result, expected)
def test_arith_coerce_scalar(data, all_arithmetic_operators):
op = tm.get_op_from_name(all_arithmetic_operators)
s = pd.Series(data)
other = 0.01
result = op(s, other)
expected = op(s.astype(float), other)
expected = expected.astype("Float64")
# rmod results in NaN that wasn't NA in original nullable Series -> unmask it
if all_arithmetic_operators == "__rmod__":
mask = (s == 0).fillna(False).to_numpy(bool)
expected.array._mask[mask] = False
tm.assert_series_equal(result, expected)
def test_array_scalar_like_equivalence(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
scalar_array = pd.array([scalar] * len(data), dtype=data.dtype)
# TODO also add len-1 array (np.array([scalar], dtype=data.dtype.numpy_dtype))
for scalar in [scalar, data.dtype.type(scalar)]:
result = op(data, scalar)
expected = op(data, scalar_array)
if isinstance(expected, ExtensionArray):
tm.assert_extension_array_equal(result, expected)
else:
# TODO div still gives float ndarray -> remove this once we have Float EA
tm.assert_numpy_array_equal(result, expected)
def test_array_NA(data, all_arithmetic_operators):
data, _ = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
scalar = pd.NA
scalar_array = pd.array([pd.NA] * len(data), dtype=data.dtype)
mask = data._mask.copy()
result = op(data, scalar)
# GH#45421 check op doesn't alter data._mask inplace
tm.assert_numpy_array_equal(mask, data._mask)
expected = op(data, scalar_array)
tm.assert_numpy_array_equal(mask, data._mask)
tm.assert_extension_array_equal(result, expected)
def test_array_scalar_like_equivalence(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
scalar_array = pd.array([scalar] * len(data), dtype=data.dtype)
# TODO also add len-1 array (np.array([scalar], dtype=data.dtype.numpy_dtype))
for scalar in [scalar, data.dtype.type(scalar)]:
if is_bool_not_implemented(data, all_arithmetic_operators):
msg = "operator '.*' not implemented for bool dtypes"
with pytest.raises(NotImplementedError, match=msg):
op(data, scalar)
with pytest.raises(NotImplementedError, match=msg):
op(data, scalar_array)
else:
result = op(data, scalar)
expected = op(data, scalar_array)
tm.assert_extension_array_equal(result, expected)
def test_frame(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
# DataFrame with scalar
df = pd.DataFrame({"A": data})
if is_bool_not_implemented(data, all_arithmetic_operators):
msg = "operator '.*' not implemented for bool dtypes"
with pytest.raises(NotImplementedError, match=msg):
op(df, scalar)
with pytest.raises(NotImplementedError, match=msg):
op(data, scalar)
return
result = op(df, scalar)
expected = pd.DataFrame({"A": op(data, scalar)})
tm.assert_frame_equal(result, expected)
def test_numpy_array_equivalence(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
numpy_array = np.array([scalar] * len(data), dtype=data.dtype.numpy_dtype)
pd_array = pd.array(numpy_array, dtype=data.dtype)
if is_bool_not_implemented(data, all_arithmetic_operators):
msg = "operator '.*' not implemented for bool dtypes"
with pytest.raises(NotImplementedError, match=msg):
op(data, numpy_array)
with pytest.raises(NotImplementedError, match=msg):
op(data, pd_array)
return
result = op(data, numpy_array)
expected = op(data, pd_array)
tm.assert_extension_array_equal(result, expected)
def test_arith_coerce_scalar(data, all_arithmetic_operators):
op = tm.get_op_from_name(all_arithmetic_operators)
s = pd.Series(data)
other = 0.01
result = op(s, other)
expected = op(s.astype(float), other)
expected = expected.astype("Float64")
# rfloordiv results in nan instead of inf
if all_arithmetic_operators == "__rfloordiv__" and np_version_under1p20:
# for numpy 1.20 https://github.com/numpy/numpy/pull/16161
# updated floordiv, now matches our behavior defined in core.ops
mask = (((expected == np.inf) |
(expected == -np.inf)).fillna(False).to_numpy(bool))
expected.array._data[mask] = np.nan
# rmod results in NaN that wasn't NA in original nullable Series -> unmask it
elif all_arithmetic_operators == "__rmod__":
mask = (s == 0).fillna(False).to_numpy(bool)
expected.array._mask[mask] = False
tm.assert_series_equal(result, expected)
def test_series(data, all_arithmetic_operators):
data, scalar = data
op = tm.get_op_from_name(all_arithmetic_operators)
check_skip(data, all_arithmetic_operators)
ser = pd.Series(data)
others = [
scalar,
np.array([scalar] * len(data), dtype=data.dtype.numpy_dtype),
pd.array([scalar] * len(data), dtype=data.dtype),
pd.Series([scalar] * len(data), dtype=data.dtype),
]
for other in others:
if is_bool_not_implemented(data, all_arithmetic_operators):
msg = "operator '.*' not implemented for bool dtypes"
with pytest.raises(NotImplementedError, match=msg):
op(ser, other)
else:
result = op(ser, other)
expected = pd.Series(op(data, other))
tm.assert_series_equal(result, expected)
def get_op_from_name(self, op_name):
return tm.get_op_from_name(op_name)
