Implement _most_ of the EA interface for DTA/TDA

pandas/core/arrays/datetimelike.py

@@ -39,7 +39,7 @@
 from pandas.core.dtypes.missing import isna
 
 import pandas.core.common as com
-from pandas.core.algorithms import checked_add_with_arr
+from pandas.core.algorithms import checked_add_with_arr, take, unique1d
 
 from .base import ExtensionOpsMixin
 from pandas.util._decorators import deprecate_kwarg
@@ -196,6 +196,71 @@ def astype(self, dtype, copy=True):
             return self._box_values(self.asi8)
         return super(DatetimeLikeArrayMixin, self).astype(dtype, copy)
 
+    # ------------------------------------------------------------------
+    # ExtensionArray Interface
+    # TODO:
+    #   * _from_sequence
+    #   * argsort / _values_for_argsort
+    #   * _reduce
+
+    def unique(self):
+        result = unique1d(self.asi8)
+        return type(self)(result, dtype=self.dtype)
+
+    def _validate_fill_value(self, fill_value):
+        """
+        If a fill_value is passed to `take` convert it to an i8 representation,
+        raising ValueError if this is not possible.
+
+        Parameters
+        ----------
+        fill_value : object
+
+        Returns
+        -------
+        fill_value : np.int64
+
+        Raises
+        ------
+        ValueError
+        """
+        raise AbstractMethodError(self)
+
+    def take(self, indices, allow_fill=False, fill_value=None):
+        if allow_fill:
+            fill_value = self._validate_fill_value(fill_value)
+
+        new_values = take(self.asi8,
+                          indices,
+                          allow_fill=allow_fill,
+                          fill_value=fill_value)
+
+        return type(self)(new_values, dtype=self.dtype)
+
+    @classmethod
+    def _concat_same_type(cls, to_concat):
+        # for TimedeltaArray and PeriodArray; DatetimeArray overrides
+        freqs = {x.freq for x in to_concat}
+        assert len(freqs) == 1
+        freq = list(freqs)[0]
+        values = np.concatenate([x.asi8 for x in to_concat])
+        return cls._simple_new(values, freq=freq)
+
+    def copy(self, deep=False):
+        values = self.asi8
+        if deep:
+            values = values.copy()
+        return type(self)(values, dtype=self.dtype, freq=self.freq)
+
+    def _values_for_factorize(self):
+        return self.asi8, iNaT
+
+    @classmethod
+    def _from_factorized(cls, values, original):
+        if is_datetime64tz_dtype(original):
+            return cls(values, tz=original.tz, freq=original.freq)
+        return cls(values, freq=original.freq)
+
     # ------------------------------------------------------------------
     # Null Handling
 

pandas/core/arrays/datetimes.py

@@ -12,7 +12,7 @@
     conversion, fields, timezones,
     resolution as libresolution)
 
-from pandas.util._decorators import cache_readonly
+from pandas.util._decorators import cache_readonly, Appender
 from pandas.errors import PerformanceWarning
 from pandas import compat
 
@@ -294,7 +294,7 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
 
             if tz is not None and index.tz is None:
                 arr = conversion.tz_localize_to_utc(
-                    ensure_int64(index.values),
+                    index.asi8,
                     tz, ambiguous=ambiguous)
 
                 index = cls(arr)
@@ -317,7 +317,7 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
         if not right_closed and len(index) and index[-1] == end:
             index = index[:-1]
 
-        return cls._simple_new(index.values, freq=freq, tz=tz)
+        return cls._simple_new(index.asi8, freq=freq, tz=tz)
 
     # -----------------------------------------------------------------
     # Descriptive Properties
@@ -419,6 +419,38 @@ def __iter__(self):
             for v in converted:
                 yield v
 
+    # ----------------------------------------------------------------
+    # ExtensionArray Interface
+
+    @property
+    def _ndarray_values(self):
+        return self._data
+
+    @Appender(dtl.DatetimeLikeArrayMixin._validate_fill_value.__doc__)
+    def _validate_fill_value(self, fill_value):
+        if isna(fill_value):
+            fill_value = iNaT
+        elif isinstance(fill_value, (datetime, np.datetime64)):
+            self._assert_tzawareness_compat(fill_value)
+            fill_value = Timestamp(fill_value).value
+        else:
+            raise ValueError("'fill_value' should be a Timestamp. "
+                             "Got '{got}'.".format(got=fill_value))
+        return fill_value
+
+    @classmethod
+    def _concat_same_type(cls, to_concat):
+        freqs = {x.freq for x in to_concat}
+        assert len(freqs) == 1
+        freq = list(freqs)[0]
+
+        tzs = {x.tz for x in to_concat}
+        assert len(tzs) == 1
+        tz = list(tzs)[0]
+
+        values = np.concatenate([x.asi8 for x in to_concat])
+        return cls._simple_new(values, freq=freq, tz=tz)
+
     # -----------------------------------------------------------------
     # Comparison Methods
 

pandas/core/arrays/period.py

@@ -216,14 +216,6 @@ def _from_sequence(cls, scalars, dtype=None, copy=False):
         ordinals = libperiod.extract_ordinals(periods, freq)
         return cls(ordinals, freq=freq)
 
-    def _values_for_factorize(self):
-        return self.asi8, iNaT
-
-    @classmethod
-    def _from_factorized(cls, values, original):
-        # type: (Sequence[Optional[Period]], PeriodArray) -> PeriodArray
-        return cls(values, freq=original.freq)
-
     @classmethod
     def _from_datetime64(cls, data, freq, tz=None):
         """Construct a PeriodArray from a datetime64 array
@@ -262,14 +254,6 @@ def _generate_range(cls, start, end, periods, freq, fields):
 
         return subarr, freq
 
-    @classmethod
-    def _concat_same_type(cls, to_concat):
-        freq = {x.freq for x in to_concat}
-        assert len(freq) == 1
-        freq = list(freq)[0]
-        values = np.concatenate([x._data for x in to_concat])
-        return cls(values, freq=freq)
-
     # --------------------------------------------------------------------
     # Data / Attributes
 
@@ -415,29 +399,20 @@ def __setitem__(
             raise TypeError(msg)
         self._data[key] = value
 
-    def take(self, indices, allow_fill=False, fill_value=None):
-        if allow_fill:
-            if isna(fill_value):
-                fill_value = iNaT
-            elif isinstance(fill_value, Period):
-                if self.freq != fill_value.freq:
-                    msg = DIFFERENT_FREQ_INDEX.format(
-                        self.freq.freqstr,
-                        fill_value.freqstr
-                    )
-                    raise IncompatibleFrequency(msg)
-
-                fill_value = fill_value.ordinal
-            else:
-                msg = "'fill_value' should be a Period. Got '{}'."
-                raise ValueError(msg.format(fill_value))
-
-        new_values = algos.take(self._data,
-                                indices,
-                                allow_fill=allow_fill,
-                                fill_value=fill_value)
-
-        return type(self)(new_values, self.freq)
+    @Appender(dtl.DatetimeLikeArrayMixin._validate_fill_value.__doc__)
+    def _validate_fill_value(self, fill_value):
+        if isna(fill_value):
+            fill_value = iNaT
+        elif isinstance(fill_value, Period):
+            if fill_value.freq != self.freq:
+                msg = DIFFERENT_FREQ_INDEX.format(self.freq.freqstr,
+                                                  fill_value.freqstr)
+                raise IncompatibleFrequency(msg)
+            fill_value = fill_value.ordinal
+        else:
+            raise ValueError("'fill_value' should be a Period. "
+                             "Got '{got}'.".format(got=fill_value))
+        return fill_value
 
     def fillna(self, value=None, method=None, limit=None):
         # TODO(#20300)

pandas/core/arrays/timedeltas.py

@@ -9,6 +9,7 @@
 from pandas._libs.tslibs.fields import get_timedelta_field
 from pandas._libs.tslibs.timedeltas import (
     array_to_timedelta64, parse_timedelta_unit)
+from pandas.util._decorators import Appender
 
 from pandas import compat
 
@@ -139,7 +140,7 @@ def _simple_new(cls, values, freq=None, dtype=_TD_DTYPE):
         result._freq = freq
         return result
 
-    def __new__(cls, values, freq=None):
+    def __new__(cls, values, freq=None, dtype=_TD_DTYPE):
 
         freq, freq_infer = dtl.maybe_infer_freq(freq)
 
@@ -193,6 +194,17 @@ def _generate_range(cls, start, end, periods, freq, closed=None):
     # ----------------------------------------------------------------
     # Array-Like / EA-Interface Methods
 
+    @Appender(dtl.DatetimeLikeArrayMixin._validate_fill_value.__doc__)
+    def _validate_fill_value(self, fill_value):
+        if isna(fill_value):
+            fill_value = iNaT
+        elif isinstance(fill_value, (timedelta, np.timedelta64, Tick)):
+            fill_value = Timedelta(fill_value).value
+        else:
+            raise ValueError("'fill_value' should be a Timedelta. "
+                             "Got '{got}'.".format(got=fill_value))
+        return fill_value
+
     # ----------------------------------------------------------------
     # Arithmetic Methods
 

pandas/core/indexes/datetimelike.py

@@ -214,6 +214,11 @@ def ceil(self, freq, ambiguous='raise', nonexistent='raise'):
 class DatetimeIndexOpsMixin(DatetimeLikeArrayMixin):
     """ common ops mixin to support a unified interface datetimelike Index """
 
+    # override DatetimeLikeArrayMixin method
+    copy = Index.copy
+    unique = Index.unique
+    take = Index.take
+
     # DatetimeLikeArrayMixin assumes subclasses are mutable, so these are
     # properties there.  They can be made into cache_readonly for Index
     # subclasses bc they are immutable
@@ -760,7 +765,7 @@ def _ensure_datetimelike_to_i8(other, to_utc=False):
         try:
             return np.array(other, copy=False).view('i8')
         except TypeError:
-            # period array cannot be coerces to int
+            # period array cannot be coerced to int
             other = Index(other)
     return other.asi8
 

pandas/core/indexes/datetimes.py

@@ -551,16 +551,12 @@ def snap(self, freq='S'):
         # TODO: what about self.name?  if so, use shallow_copy?
 
     def unique(self, level=None):
-        # Override here since IndexOpsMixin.unique uses self._values.unique
-        # For DatetimeIndex with TZ, that's a DatetimeIndex -> recursion error
-        # So we extract the tz-naive DatetimeIndex, unique that, and wrap the
-        # result with out TZ.
-        if self.tz is not None:
-            naive = type(self)(self._ndarray_values, copy=False)
-        else:
-            naive = self
-        result = super(DatetimeIndex, naive).unique(level=level)
-        return self._shallow_copy(result.values)
+        if level is not None:
+            self._validate_index_level(level)
+
+        # call DatetimeArray method
+        result = DatetimeArrayMixin.unique(self)
+        return self._shallow_copy(result._data)
 
     def union(self, other):
         """

pandas/tests/arrays/test_datetimelike.py

@@ -175,6 +175,41 @@ def test_int_properties(self, datetime_index, propname):
 
         tm.assert_numpy_array_equal(result, expected)
 
+    def test_take_fill_valid(self, datetime_index, tz_naive_fixture):
+        dti = datetime_index.tz_localize(tz_naive_fixture)
+        arr = DatetimeArray(dti)
+
+        now = pd.Timestamp.now().tz_localize(dti.tz)
+        result = arr.take([-1, 1], allow_fill=True, fill_value=now)
+        assert result[0] == now
+
+        with pytest.raises(ValueError):
+            # fill_value Timedelta invalid
+            arr.take([-1, 1], allow_fill=True, fill_value=now - now)
+
+        with pytest.raises(ValueError):
+            # fill_value Period invalid
+            arr.take([-1, 1], allow_fill=True, fill_value=pd.Period('2014Q1'))
+
+        tz = None if dti.tz is not None else 'US/Eastern'
+        now = pd.Timestamp.now().tz_localize(tz)
+        with pytest.raises(TypeError):
+            # Timestamp with mismatched tz-awareness
+            arr.take([-1, 1], allow_fill=True, fill_value=now)
+
+    def test_concat_same_type_invalid(self, datetime_index):
+        # different timezones
+        dti = datetime_index
+        arr = DatetimeArray(dti)
+
+        if arr.tz is None:
+            other = arr.tz_localize('UTC')
+        else:
+            other = arr.tz_localize(None)
+
+        with pytest.raises(AssertionError):
+            arr._concat_same_type([arr, other])
+
 
 class TestTimedeltaArray(object):
     def test_from_tdi(self):
@@ -223,6 +258,35 @@ def test_int_properties(self, timedelta_index, propname):
 
         tm.assert_numpy_array_equal(result, expected)
 
+    def test_take_fill_valid(self, timedelta_index):
+        tdi = timedelta_index
+        arr = TimedeltaArray(tdi)
+
+        td1 = pd.Timedelta(days=1)
+        result = arr.take([-1, 1], allow_fill=True, fill_value=td1)
+        assert result[0] == td1
+
+        now = pd.Timestamp.now()
+        with pytest.raises(ValueError):
+            # fill_value Timestamp invalid
+            arr.take([0, 1], allow_fill=True, fill_value=now)
+
+        with pytest.raises(ValueError):
+            # fill_value Period invalid
+            arr.take([0, 1], allow_fill=True, fill_value=now.to_period('D'))
+
+    def test_concat_same_type_invalid(self, timedelta_index):
+        # different freqs
+        tdi = timedelta_index
+        arr = TimedeltaArray(tdi)
+        other = pd.timedelta_range('1D', periods=5, freq='2D')
+        # FIXME: TimedeltaArray should inherit freq='2D' without specifying it
+        other = TimedeltaArray(other, freq='2D')
+        assert other.freq != arr.freq
+
+        with pytest.raises(AssertionError):
+            arr._concat_same_type([arr, other])
+
 
 class TestPeriodArray(object):