Skip to content

Commit df51656

Browse files
authored
TST: fix tm.to_array nullable handling (#44884)
1 parent eabac61 commit df51656

File tree

2 files changed

+14
-30
lines changed

2 files changed

+14
-30
lines changed

pandas/_testing/__init__.py

Lines changed: 8 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,12 @@
2828
from pandas._typing import Dtype
2929

3030
from pandas.core.dtypes.common import (
31-
is_datetime64_dtype,
32-
is_datetime64tz_dtype,
3331
is_float_dtype,
3432
is_integer_dtype,
35-
is_period_dtype,
3633
is_sequence,
37-
is_timedelta64_dtype,
3834
is_unsigned_integer_dtype,
3935
pandas_dtype,
4036
)
41-
from pandas.core.dtypes.dtypes import IntervalDtype
4237

4338
import pandas as pd
4439
from pandas import (
@@ -112,14 +107,11 @@
112107
)
113108
from pandas.core.arrays import (
114109
BaseMaskedArray,
115-
DatetimeArray,
116110
ExtensionArray,
117111
PandasArray,
118-
PeriodArray,
119-
TimedeltaArray,
120-
period_array,
121112
)
122113
from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
114+
from pandas.core.construction import extract_array
123115

124116
if TYPE_CHECKING:
125117
from pandas import (
@@ -268,13 +260,6 @@ def box_expected(expected, box_cls, transpose=True):
268260
# single-row special cases in datetime arithmetic
269261
expected = expected.T
270262
expected = pd.concat([expected] * 2, ignore_index=True)
271-
elif box_cls is PeriodArray:
272-
# the PeriodArray constructor is not as flexible as period_array
273-
expected = period_array(expected)
274-
elif box_cls is DatetimeArray:
275-
expected = DatetimeArray(expected)
276-
elif box_cls is TimedeltaArray:
277-
expected = TimedeltaArray(expected)
278263
elif box_cls is np.ndarray or box_cls is np.array:
279264
expected = np.array(expected)
280265
elif box_cls is to_array:
@@ -285,21 +270,16 @@ def box_expected(expected, box_cls, transpose=True):
285270

286271

287272
def to_array(obj):
273+
"""
274+
Similar to pd.array, but does not cast numpy dtypes to nullable dtypes.
275+
"""
288276
# temporary implementation until we get pd.array in place
289277
dtype = getattr(obj, "dtype", None)
290278

291-
if is_period_dtype(dtype):
292-
return period_array(obj)
293-
elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
294-
return DatetimeArray._from_sequence(obj)
295-
elif is_timedelta64_dtype(dtype):
296-
return TimedeltaArray._from_sequence(obj)
297-
elif isinstance(obj, pd.core.arrays.BooleanArray):
298-
return obj
299-
elif isinstance(dtype, IntervalDtype):
300-
return pd.core.arrays.IntervalArray(obj)
301-
else:
302-
return np.array(obj)
279+
if dtype is None:
280+
return np.asarray(obj)
281+
282+
return extract_array(obj, extract_numpy=True)
303283

304284

305285
# -----------------------------------------------------------------------------

pandas/tests/arithmetic/test_numeric.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1400,11 +1400,15 @@ def test_integer_array_add_list_like(
14001400
if Series == box_pandas_1d_array:
14011401
expected = Series(expected_data, dtype="Int64")
14021402
elif Series == box_1d_array:
1403-
expected = Series(expected_data, dtype="object")
1403+
if box_pandas_1d_array is tm.to_array:
1404+
expected = Series(expected_data, dtype="Int64")
1405+
else:
1406+
expected = Series(expected_data, dtype="object")
14041407
elif Index in (box_pandas_1d_array, box_1d_array):
14051408
expected = Int64Index(expected_data)
14061409
else:
1407-
expected = np.array(expected_data, dtype="object")
1410+
# box_pandas_1d_array is tm.to_array; preserves IntegerArray
1411+
expected = array(expected_data, dtype="Int64")
14081412

14091413
tm.assert_equal(left, expected)
14101414
tm.assert_equal(right, expected)

0 commit comments

Comments
 (0)