implement constructors for TimedeltaArray, DatetimeArray (pandas-dev#21803)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -18,7 +18,37 @@
 from pandas.core.algorithms import checked_add_with_arr
 
 
-class DatetimeLikeArrayMixin(object):
+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):
+        """return an attributes dict for my class"""
+        return {k: getattr(self, k, None) for k in self._attributes}
+
+    def _shallow_copy(self, values=None, **kwargs):
+        if values is None:
+            # Note: slightly different from Index implementation which defaults
+            # to self.values
+            values = self._ndarray_values
+
+        attributes = self._get_attributes_dict()
+        attributes.update(kwargs)
+        if not len(values) and 'dtype' not in kwargs:
+            attributes['dtype'] = self.dtype
+        return self._simple_new(values, **attributes)
+
+
+class DatetimeLikeArrayMixin(AttributesMixin):
     """
     Shared Base/Mixin class for DatetimeArray, TimedeltaArray, PeriodArray
 
@@ -56,9 +86,61 @@ def asi8(self):
         # do not cache or you'll create a memory leak
         return self.values.view('i8')
 
+    # ------------------------------------------------------------------
+    # Array-like Methods
+
     def __len__(self):
         return len(self._data)
 
+    def __getitem__(self, key):
+        """
+        This getitem defers to the underlying array, which by-definition can
+        only handle list-likes, slices, and integer scalars
+        """
+
+        is_int = lib.is_integer(key)
+        if lib.is_scalar(key) and not is_int:
+            raise IndexError("only integers, slices (`:`), ellipsis (`...`), "
+                             "numpy.newaxis (`None`) and integer or boolean "
+                             "arrays are valid indices")
+
+        getitem = self._data.__getitem__
+        if is_int:
+            val = getitem(key)
+            return self._box_func(val)
+        else:
+            if com.is_bool_indexer(key):
+                key = np.asarray(key)
+                if key.all():
+                    key = slice(0, None, None)
+                else:
+                    key = lib.maybe_booleans_to_slice(key.view(np.uint8))
+
+            attribs = self._get_attributes_dict()
+
+            is_period = is_period_dtype(self)
+            if is_period:
+                freq = self.freq
+            else:
+                freq = None
+                if isinstance(key, slice):
+                    if self.freq is not None and key.step is not None:
+                        freq = key.step * self.freq
+                    else:
+                        freq = self.freq
+
+            attribs['freq'] = freq
+
+            result = getitem(key)
+            if result.ndim > 1:
+                # To support MPL which performs slicing with 2 dim
+                # even though it only has 1 dim by definition
+                if is_period:
+                    return self._simple_new(result, **attribs)
+                return result
+
+            return self._simple_new(result, **attribs)
+
     # ------------------------------------------------------------------
     # Null Handling
 
@@ -97,6 +179,27 @@ def _maybe_mask_results(self, result, fill_value=None, convert=None):
             result[self._isnan] = fill_value
         return result
 
+    def _nat_new(self, box=True):
+        """
+        Return Array/Index or ndarray filled with NaT which has the same
+        length as the caller.
+
+        Parameters
+        ----------
+        box : boolean, default True
+            - If True returns a Array/Index as the same as caller.
+            - If False returns ndarray of np.int64.
+        """
+        result = np.zeros(len(self), dtype=np.int64)
+        result.fill(iNaT)
+        if not box:
+            return result
+
+        attribs = self._get_attributes_dict()
+        if not is_period_dtype(self):
+            attribs['freq'] = None
+        return self._simple_new(result, **attribs)
+
     # ------------------------------------------------------------------
     # Frequency Properties/Methods
 
@@ -195,6 +298,17 @@ def _add_delta_tdi(self, other):
             new_values[mask] = iNaT
         return new_values.view('i8')
 
+    def _add_nat(self):
+        """Add pd.NaT to self"""
+        if is_period_dtype(self):
+            raise TypeError('Cannot add {cls} and {typ}'
+                            .format(cls=type(self).__name__,
+                                    typ=type(NaT).__name__))
+
+        # GH#19124 pd.NaT is treated like a timedelta for both timedelta
+        # and datetime dtypes
+        return self._nat_new(box=True)
+
     def _sub_nat(self):
         """Subtract pd.NaT from self"""
         # GH#19124 Timedelta - datetime is not in general well-defined.