Implement take for EA mixins

pandas/core/arrays/datetimelike.py

@@ -11,7 +11,8 @@
 from pandas._libs.tslibs.period import (
     Period, DIFFERENT_FREQ_INDEX, IncompatibleFrequency)
 
-from pandas.errors import NullFrequencyError, PerformanceWarning
+from pandas.errors import (
+    NullFrequencyError, PerformanceWarning, AbstractMethodError)
 from pandas import compat
 
 from pandas.tseries import frequencies
@@ -36,7 +37,7 @@
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
 
 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
 
 from .base import ExtensionOpsMixin
 from pandas.util._decorators import deprecate_kwarg
@@ -77,12 +78,10 @@ class AttributesMixin(object):
     @property
     def _attributes(self):
         # Inheriting subclass should implement _attributes as a list of strings
-        from pandas.errors import AbstractMethodError
         raise AbstractMethodError(self)
 
     @classmethod
     def _simple_new(cls, values, **kwargs):
-        from pandas.errors import AbstractMethodError
         raise AbstractMethodError(cls)
 
     def _get_attributes_dict(self):
@@ -119,7 +118,7 @@ def _box_func(self):
         """
         box function to get object from internal representation
         """
-        raise com.AbstractMethodError(self)
+        raise AbstractMethodError(self)
 
     def _box_values(self, values):
         """
@@ -140,6 +139,67 @@ def asi8(self):
         # do not cache or you'll create a memory leak
         return self.values.view('i8')
 
+    # ------------------------------------------------------------------
+    # Extension Array Interface
+    # TODO:
+    #   _from_sequence
+    #   _from_factorized
+    #   __setitem__
+    #   _values_for_argsort
+    #   argsort
+    #   fillna
+    #   dropna
+    #   shift
+    #   unique
+    #   _values_for_factorize
+    #   factorize
+    #   _formatting_values
+    #   _reduce
+    #   copy
+
+    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._data,
+                          indices,
+                          allow_fill=allow_fill,
+                          fill_value=fill_value)
+
+        # TODO: use "infer"?  Why does not passing freq cause
+        #  failures in py37 but not py27?
+        freq = self.freq if is_period_dtype(self) else None
+        return self._shallow_copy(new_values, freq=freq)
+
+    @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._data for x in to_concat])
+        return cls._simple_new(values, freq=freq)
+
     # ------------------------------------------------------------------
     # Array-like Methods
 
@@ -211,6 +271,10 @@ def astype(self, dtype, copy=True):
     # ------------------------------------------------------------------
     # Null Handling
 
+    def isna(self):
+        # EA Interface
+        return self._isnan
+
     @property  # NB: override with cache_readonly in immutable subclasses
     def _isnan(self):
         """ return if each value is nan"""
@@ -352,13 +416,13 @@ def _add_datelike(self, other):
                                 typ=type(other).__name__))
 
     def _sub_datelike(self, other):
-        raise com.AbstractMethodError(self)
+        raise AbstractMethodError(self)
 
     def _sub_period(self, other):
         return NotImplemented
 
     def _add_offset(self, offset):
-        raise com.AbstractMethodError(self)
+        raise AbstractMethodError(self)
 
     def _add_delta(self, other):
         return NotImplemented
@@ -380,7 +444,7 @@ def _add_delta_tdi(self, other):
         Add a delta of a TimedeltaIndex
         return the i8 result view
         """
-        if not len(self) == len(other):
+        if len(self) != len(other):
             raise ValueError("cannot add indices of unequal length")
 
         if isinstance(other, np.ndarray):

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
 
@@ -298,6 +298,35 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
 
         return cls._simple_new(index.values, freq=freq, tz=tz)
 
+    # ----------------------------------------------------------------
+    # Extension Array Interface
+
+    @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):
+        # for TimedeltaArray and PeriodArray; DatetimeArray requires tz
+        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._data for x in to_concat])
+        return cls._simple_new(values, freq=freq, tz=tz)
+
     # -----------------------------------------------------------------
     # Descriptive Properties
 

pandas/core/arrays/period.py

@@ -15,13 +15,14 @@
 from pandas._libs.tslibs.fields import isleapyear_arr
 
 from pandas import compat
-from pandas.util._decorators import (cache_readonly, deprecate_kwarg)
+from pandas.util._decorators import cache_readonly, deprecate_kwarg, Appender
 
 from pandas.core.dtypes.common import (
     is_integer_dtype, is_float_dtype, is_period_dtype, is_timedelta64_dtype,
     is_datetime64_dtype, _TD_DTYPE)
 from pandas.core.dtypes.dtypes import PeriodDtype
 from pandas.core.dtypes.generic import ABCSeries
+from pandas.core.dtypes.missing import isna
 
 import pandas.core.common as com
 
@@ -193,6 +194,23 @@ def _generate_range(cls, start, end, periods, freq, fields):
 
         return subarr, freq
 
+    # --------------------------------------------------------------------
+    # ExtensionArray Interface
+
+    @Appender(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:
+                raise ValueError("'fill_value' freq must match own "
+                                 "freq ({freq})".format(freq=self.freq))
+            fill_value = fill_value.ordinal
+        else:
+            raise ValueError("'fill_value' should be a Period. "
+                             "Got '{got}'.".format(got=fill_value))
+        return fill_value
+
     # --------------------------------------------------------------------
     # Vectorized analogues of Period properties
 

pandas/core/arrays/timedeltas.py

@@ -8,6 +8,8 @@
 from pandas._libs.tslibs.fields import get_timedelta_field
 from pandas._libs.tslibs.timedeltas import array_to_timedelta64
 
+from pandas.util._decorators import Appender
+
 from pandas import compat
 
 from pandas.core.dtypes.common import (
@@ -180,6 +182,20 @@ def _generate_range(cls, start, end, periods, freq, closed=None, **kwargs):
 
         return index
 
+    # ----------------------------------------------------------------
+    # Extension Array Interface
+
+    @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