From 95577a3ef830f9a5c85071b58971dbb95c651117 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 14 Dec 2021 11:17:06 -0800 Subject: [PATCH] TST: fix tm.to_array nullable handling --- pandas/_testing/__init__.py | 36 ++++++------------------- pandas/tests/arithmetic/test_numeric.py | 8 ++++-- 2 files changed, 14 insertions(+), 30 deletions(-) diff --git a/pandas/_testing/__init__.py b/pandas/_testing/__init__.py index 16094bd88d66f..70f7deff38d94 100644 --- a/pandas/_testing/__init__.py +++ b/pandas/_testing/__init__.py @@ -28,17 +28,12 @@ from pandas._typing import Dtype from pandas.core.dtypes.common import ( - is_datetime64_dtype, - is_datetime64tz_dtype, is_float_dtype, is_integer_dtype, - is_period_dtype, is_sequence, - is_timedelta64_dtype, is_unsigned_integer_dtype, pandas_dtype, ) -from pandas.core.dtypes.dtypes import IntervalDtype import pandas as pd from pandas import ( @@ -112,14 +107,11 @@ ) from pandas.core.arrays import ( BaseMaskedArray, - DatetimeArray, ExtensionArray, PandasArray, - PeriodArray, - TimedeltaArray, - period_array, ) from pandas.core.arrays._mixins import NDArrayBackedExtensionArray +from pandas.core.construction import extract_array if TYPE_CHECKING: from pandas import ( @@ -257,13 +249,6 @@ def box_expected(expected, box_cls, transpose=True): # single-row special cases in datetime arithmetic expected = expected.T expected = pd.concat([expected] * 2, ignore_index=True) - elif box_cls is PeriodArray: - # the PeriodArray constructor is not as flexible as period_array - expected = period_array(expected) - elif box_cls is DatetimeArray: - expected = DatetimeArray(expected) - elif box_cls is TimedeltaArray: - expected = TimedeltaArray(expected) elif box_cls is np.ndarray or box_cls is np.array: expected = np.array(expected) elif box_cls is to_array: @@ -274,21 +259,16 @@ def box_expected(expected, box_cls, transpose=True): def to_array(obj): + """ + Similar to pd.array, but does not cast numpy dtypes to nullable dtypes. + """ # temporary implementation until we get pd.array in place dtype = getattr(obj, "dtype", None) - if is_period_dtype(dtype): - return period_array(obj) - elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype): - return DatetimeArray._from_sequence(obj) - elif is_timedelta64_dtype(dtype): - return TimedeltaArray._from_sequence(obj) - elif isinstance(obj, pd.core.arrays.BooleanArray): - return obj - elif isinstance(dtype, IntervalDtype): - return pd.core.arrays.IntervalArray(obj) - else: - return np.array(obj) + if dtype is None: + return np.asarray(obj) + + return extract_array(obj, extract_numpy=True) # ----------------------------------------------------------------------------- diff --git a/pandas/tests/arithmetic/test_numeric.py b/pandas/tests/arithmetic/test_numeric.py index 4935151ec986e..2cc49d38471c3 100644 --- a/pandas/tests/arithmetic/test_numeric.py +++ b/pandas/tests/arithmetic/test_numeric.py @@ -1400,11 +1400,15 @@ def test_integer_array_add_list_like( if Series == box_pandas_1d_array: expected = Series(expected_data, dtype="Int64") elif Series == box_1d_array: - expected = Series(expected_data, dtype="object") + if box_pandas_1d_array is tm.to_array: + expected = Series(expected_data, dtype="Int64") + else: + expected = Series(expected_data, dtype="object") elif Index in (box_pandas_1d_array, box_1d_array): expected = Int64Index(expected_data) else: - expected = np.array(expected_data, dtype="object") + # box_pandas_1d_array is tm.to_array; preserves IntegerArray + expected = array(expected_data, dtype="Int64") tm.assert_equal(left, expected) tm.assert_equal(right, expected)