72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from pandas.core.dtypes.cast import maybe_upcast_putmask
|
|
|
|
from pandas import Series
|
|
import pandas._testing as tm
|
|
|
|
|
|
@pytest.mark.parametrize("result", [Series([10, 11, 12]), [10, 11, 12], (10, 11, 12)])
|
|
def test_upcast_error(result):
|
|
# GH23823 require result arg to be ndarray
|
|
mask = np.array([False, True, False])
|
|
other = np.array([61, 62, 63])
|
|
with pytest.raises(ValueError, match="The result input must be a ndarray"):
|
|
result, _ = maybe_upcast_putmask(result, mask, other)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"arr, other",
|
|
[
|
|
(np.arange(1, 6), np.array([61, 62, 63])),
|
|
(np.arange(1, 6), np.array([61.1, 62.2, 63.3])),
|
|
(np.arange(10, 15), np.array([61, 62])),
|
|
(np.arange(10, 15), np.array([61, np.nan])),
|
|
(
|
|
np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]"),
|
|
np.arange("2018-01-01", "2018-01-04", dtype="datetime64[D]"),
|
|
),
|
|
(
|
|
np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]"),
|
|
np.arange("2018-01-01", "2018-01-03", dtype="datetime64[D]"),
|
|
),
|
|
],
|
|
)
|
|
def test_upcast_scalar_other(arr, other):
|
|
# for now we do not support non-scalar `other`
|
|
mask = np.array([False, True, False, True, True])
|
|
with pytest.raises(ValueError, match="other must be a scalar"):
|
|
maybe_upcast_putmask(arr, mask, other)
|
|
|
|
|
|
def test_upcast():
|
|
# GH23823
|
|
arr = np.arange(1, 6)
|
|
mask = np.array([False, True, False, True, True])
|
|
result, changed = maybe_upcast_putmask(arr, mask, other=np.nan)
|
|
|
|
expected = np.array([1, np.nan, 3, np.nan, np.nan])
|
|
assert changed
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
|
|
def test_upcast_datetime():
|
|
# GH23823
|
|
arr = np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]")
|
|
mask = np.array([False, True, False, True, True])
|
|
result, changed = maybe_upcast_putmask(arr, mask, other=np.nan)
|
|
|
|
expected = np.array(
|
|
[
|
|
"2019-01-01",
|
|
np.datetime64("NaT"),
|
|
"2019-01-03",
|
|
np.datetime64("NaT"),
|
|
np.datetime64("NaT"),
|
|
],
|
|
dtype="datetime64[D]",
|
|
)
|
|
assert not changed
|
|
tm.assert_numpy_array_equal(result, expected)
|