Skip to content

Commit

Permalink
API: consistent __array__ for datetime-like ExtensionArrays (pandas-d…
Browse files Browse the repository at this point in the history
  • Loading branch information
jorisvandenbossche authored and Pingviinituutti committed Feb 28, 2019
1 parent 9c6f682 commit deef8e4
Show file tree
Hide file tree
Showing 6 changed files with 139 additions and 23 deletions.
6 changes: 6 additions & 0 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,12 @@ def _formatter(self, boxed=False):
def nbytes(self):
return self._data.nbytes

def __array__(self, dtype=None):
# used for Timedelta/DatetimeArray, overwritten by PeriodArray
if is_object_dtype(dtype):
return np.array(list(self), dtype=object)
return self._data

@property
def shape(self):
return (len(self),)
Expand Down
13 changes: 6 additions & 7 deletions pandas/core/arrays/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@
from pandas.core.dtypes.common import (
_INT64_DTYPE, _NS_DTYPE, is_categorical_dtype, is_datetime64_dtype,
is_datetime64_ns_dtype, is_datetime64tz_dtype, is_dtype_equal,
is_extension_type, is_float_dtype, is_int64_dtype, is_object_dtype,
is_period_dtype, is_string_dtype, is_timedelta64_dtype, pandas_dtype)
is_extension_type, is_float_dtype, is_object_dtype, is_period_dtype,
is_string_dtype, is_timedelta64_dtype, pandas_dtype)
from pandas.core.dtypes.dtypes import DatetimeTZDtype
from pandas.core.dtypes.generic import ABCIndexClass, ABCPandasArray, ABCSeries
from pandas.core.dtypes.missing import isna
Expand Down Expand Up @@ -524,12 +524,11 @@ def _resolution(self):
# Array-Like / EA-Interface Methods

def __array__(self, dtype=None):
if is_object_dtype(dtype) or (dtype is None and self.tz):
return np.array(list(self), dtype=object)
elif is_int64_dtype(dtype):
return self.asi8
if dtype is None and self.tz:
# The default for tz-aware is object, to preserve tz info
dtype = object

return self._data
return super(DatetimeArray, self).__array__(dtype=dtype)

def __iter__(self):
"""
Expand Down
4 changes: 4 additions & 0 deletions pandas/core/arrays/period.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,10 @@ def freq(self):
"""
return self.dtype.freq

def __array__(self, dtype=None):
# overriding DatetimelikeArray
return np.array(list(self), dtype=object)

# --------------------------------------------------------------------
# Vectorized analogues of Period properties

Expand Down
12 changes: 1 addition & 11 deletions pandas/core/arrays/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

from pandas.core.dtypes.common import (
_NS_DTYPE, _TD_DTYPE, ensure_int64, is_datetime64_dtype, is_float_dtype,
is_int64_dtype, is_integer_dtype, is_list_like, is_object_dtype, is_scalar,
is_integer_dtype, is_list_like, is_object_dtype, is_scalar,
is_string_dtype, is_timedelta64_dtype, is_timedelta64_ns_dtype,
pandas_dtype)
from pandas.core.dtypes.dtypes import DatetimeTZDtype
Expand Down Expand Up @@ -265,16 +265,6 @@ def _maybe_clear_freq(self):
# ----------------------------------------------------------------
# Array-Like / EA-Interface Methods

def __array__(self, dtype=None):
# TODO(https://github.com/pandas-dev/pandas/pull/23593)
# Maybe push to parent once datetimetz __array__ is figured out.
if is_object_dtype(dtype):
return np.array(list(self), dtype=object)
elif is_int64_dtype(dtype):
return self.asi8

return self._data

@Appender(dtl.DatetimeLikeArrayMixin._validate_fill_value.__doc__)
def _validate_fill_value(self, fill_value):
if isna(fill_value):
Expand Down
122 changes: 117 additions & 5 deletions pandas/tests/arrays/test_datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,48 @@ def test_round(self, tz_naive_fixture):
expected = dti - pd.Timedelta(minutes=1)
tm.assert_index_equal(result, expected)

def test_array_interface(self, datetime_index):
arr = DatetimeArray(datetime_index)

# default asarray gives the same underlying data (for tz naive)
result = np.asarray(arr)
expected = arr._data
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, copy=False)
assert result is expected
tm.assert_numpy_array_equal(result, expected)

# specifying M8[ns] gives the same result as default
result = np.asarray(arr, dtype='datetime64[ns]')
expected = arr._data
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, dtype='datetime64[ns]', copy=False)
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, dtype='datetime64[ns]')
assert result is not expected
tm.assert_numpy_array_equal(result, expected)

# to object dtype
result = np.asarray(arr, dtype=object)
expected = np.array(list(arr), dtype=object)
tm.assert_numpy_array_equal(result, expected)

# to other dtype always copies
result = np.asarray(arr, dtype='int64')
assert result is not arr.asi8
assert not np.may_share_memory(arr, result)
expected = arr.asi8.copy()
tm.assert_numpy_array_equal(result, expected)

# other dtypes handled by numpy
for dtype in ['float64', str]:
result = np.asarray(arr, dtype=dtype)
expected = np.asarray(arr).astype(dtype)
tm.assert_numpy_array_equal(result, expected)

def test_array_object_dtype(self, tz_naive_fixture):
# GH#23524
tz = tz_naive_fixture
Expand All @@ -255,7 +297,7 @@ def test_array_object_dtype(self, tz_naive_fixture):
result = np.array(dti, dtype=object)
tm.assert_numpy_array_equal(result, expected)

def test_array(self, tz_naive_fixture):
def test_array_tz(self, tz_naive_fixture):
# GH#23524
tz = tz_naive_fixture
dti = pd.date_range('2016-01-01', periods=3, tz=tz)
Expand All @@ -265,13 +307,18 @@ def test_array(self, tz_naive_fixture):
result = np.array(arr, dtype='M8[ns]')
tm.assert_numpy_array_equal(result, expected)

result = np.array(arr, dtype='datetime64[ns]')
tm.assert_numpy_array_equal(result, expected)

# check that we are not making copies when setting copy=False
result = np.array(arr, dtype='M8[ns]', copy=False)
assert result.base is expected.base
assert result.base is not None
result = np.array(arr, dtype='datetime64[ns]', copy=False)
assert result.base is expected.base
assert result.base is not None

def test_array_i8_dtype(self, tz_naive_fixture):
# GH#23524
tz = tz_naive_fixture
dti = pd.date_range('2016-01-01', periods=3, tz=tz)
arr = DatetimeArray(dti)
Expand All @@ -283,10 +330,10 @@ def test_array_i8_dtype(self, tz_naive_fixture):
result = np.array(arr, dtype=np.int64)
tm.assert_numpy_array_equal(result, expected)

# check that we are not making copies when setting copy=False
# check that we are still making copies when setting copy=False
result = np.array(arr, dtype='i8', copy=False)
assert result.base is expected.base
assert result.base is not None
assert result.base is not expected.base
assert result.base is None

def test_from_array_keeps_base(self):
# Ensure that DatetimeArray._data.base isn't lost.
Expand Down Expand Up @@ -470,6 +517,48 @@ def test_int_properties(self, timedelta_index, propname):

tm.assert_numpy_array_equal(result, expected)

def test_array_interface(self, timedelta_index):
arr = TimedeltaArray(timedelta_index)

# default asarray gives the same underlying data
result = np.asarray(arr)
expected = arr._data
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, copy=False)
assert result is expected
tm.assert_numpy_array_equal(result, expected)

# specifying m8[ns] gives the same result as default
result = np.asarray(arr, dtype='timedelta64[ns]')
expected = arr._data
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, dtype='timedelta64[ns]', copy=False)
assert result is expected
tm.assert_numpy_array_equal(result, expected)
result = np.array(arr, dtype='timedelta64[ns]')
assert result is not expected
tm.assert_numpy_array_equal(result, expected)

# to object dtype
result = np.asarray(arr, dtype=object)
expected = np.array(list(arr), dtype=object)
tm.assert_numpy_array_equal(result, expected)

# to other dtype always copies
result = np.asarray(arr, dtype='int64')
assert result is not arr.asi8
assert not np.may_share_memory(arr, result)
expected = arr.asi8.copy()
tm.assert_numpy_array_equal(result, expected)

# other dtypes handled by numpy
for dtype in ['float64', str]:
result = np.asarray(arr, dtype=dtype)
expected = np.asarray(arr).astype(dtype)
tm.assert_numpy_array_equal(result, expected)

def test_take_fill_valid(self, timedelta_index):
tdi = timedelta_index
arr = TimedeltaArray(tdi)
Expand Down Expand Up @@ -543,3 +632,26 @@ def test_int_properties(self, period_index, propname):
expected = np.array(getattr(pi, propname))

tm.assert_numpy_array_equal(result, expected)

def test_array_interface(self, period_index):
arr = PeriodArray(period_index)

# default asarray gives objects
result = np.asarray(arr)
expected = np.array(list(arr), dtype=object)
tm.assert_numpy_array_equal(result, expected)

# to object dtype (same as default)
result = np.asarray(arr, dtype=object)
tm.assert_numpy_array_equal(result, expected)

# to other dtypes
with pytest.raises(TypeError):
np.asarray(arr, dtype='int64')

with pytest.raises(TypeError):
np.asarray(arr, dtype='float64')

result = np.asarray(arr, dtype='S20')
expected = np.asarray(arr).astype('S20')
tm.assert_numpy_array_equal(result, expected)
5 changes: 5 additions & 0 deletions pandas/tests/extension/base/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pandas.core.dtypes.dtypes import ExtensionDtype

import pandas as pd
import pandas.util.testing as tm

from .base import BaseExtensionTests

Expand Down Expand Up @@ -33,6 +34,10 @@ def test_array_interface(self, data):
result = np.array(data)
assert result[0] == data[0]

result = np.array(data, dtype=object)
expected = np.array(list(data), dtype=object)
tm.assert_numpy_array_equal(result, expected)

def test_is_extension_array_dtype(self, data):
assert is_extension_array_dtype(data)
assert is_extension_array_dtype(data.dtype)
Expand Down

0 comments on commit deef8e4

Please sign in to comment.