API: consistent __array__ for datetime-like ExtensionArrays (#23593)

Author: jorisvandenbossche

pandas/core/arrays/datetimelike.py

@@ -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),)

pandas/core/arrays/datetimes.py

@@ -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
@@ -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):
         """

pandas/core/arrays/period.py

@@ -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
 

pandas/core/arrays/timedeltas.py

@@ -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
@@ -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):

pandas/tests/arrays/test_datetimelike.py

@@ -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
@@ -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)
@@ -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)
@@ -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.
@@ -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)
@@ -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)

pandas/tests/extension/base/interface.py

@@ -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
 
@@ -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)