ENH: time slicing on datetime indicies

stephenwlin · wesm · commit c834758f676a · 2013-02-10T17:32:50.000-05:00
diff --git a/pandas/core/index.py b/pandas/core/index.py
@@ -1117,6 +1117,30 @@ def _wrap_joined_index(self, joined, other):
         name = self.name if self.name == other.name else None
         return Index(joined, name=name)
 
+    def slice_indexer(self, start=None, end=None, step=None):
+        """
+        For an ordered Index, compute the slice indexer for input labels and
+        step
+
+        Parameters
+        ----------
+        start : label, default None
+            If None, defaults to the beginning
+        end : label, default None
+            If None, defaults to the end
+        step : int, default None 
+
+        Returns
+        -------
+        indexer : ndarray or slice
+
+        Notes
+        -----
+        This function assumes that the data is sorted, so use at your own peril
+        """
+        start_slice, end_slice = self.slice_locs(start, end)
+        return slice(start_slice, end_slice, step)
+
     def slice_locs(self, start=None, end=None):
         """
         For an ordered Index, compute the slice locations for input labels
@@ -1125,27 +1149,27 @@ def slice_locs(self, start=None, end=None):
         ----------
         start : label, default None
             If None, defaults to the beginning
-        end : label
+        end : label, default None
             If None, defaults to the end
 
         Returns
         -------
-        (begin, end) : (int, int)
+        (start, end) : (int, int)
 
         Notes
         -----
         This function assumes that the data is sorted, so use at your own peril
         """
         if start is None:
-            beg_slice = 0
+            start_slice = 0
         else:
             try:
-                beg_slice = self.get_loc(start)
-                if isinstance(beg_slice, slice):
-                    beg_slice = beg_slice.start
+                start_slice = self.get_loc(start)
+                if isinstance(start_slice, slice):
+                    start_slice = start_slice.start
             except KeyError:
                 if self.is_monotonic:
-                    beg_slice = self.searchsorted(start, side='left')
+                    start_slice = self.searchsorted(start, side='left')
                 else:
                     raise
 
@@ -1164,7 +1188,7 @@ def slice_locs(self, start=None, end=None):
                 else:
                     raise
 
-        return beg_slice, end_slice
+        return start_slice, end_slice
 
     def delete(self, loc):
         """
@@ -2106,7 +2130,7 @@ def slice_locs(self, start=None, end=None, strict=False):
 
         Returns
         -------
-        (begin, end) : (int, int)
+        (start, end) : (int, int)
 
         Notes
         -----
diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py
@@ -462,20 +462,19 @@ def _convert_to_indexer(self, obj, axis=0):
                     raise
 
             if null_slice or position_slice:
-                slicer = obj
+                indexer = obj
             else:
                 try:
-                    i, j = labels.slice_locs(start, stop)
-                    slicer = slice(i, j, obj.step)
+                    indexer = labels.slice_indexer(start, stop, obj.step)
                 except Exception:
                     if _is_index_slice(obj):
                         if labels.inferred_type == 'integer':
                             raise
-                        slicer = obj
+                        indexer = obj
                     else:
                         raise
 
-            return slicer
+            return indexer
 
         elif _is_list_like(obj):
             if com._is_bool_indexer(obj):
@@ -535,8 +534,10 @@ def _tuplify(self, loc):
     def _get_slice_axis(self, slice_obj, axis=0):
         obj = self.obj
 
-        axis_name = obj._get_axis_name(axis)
-        labels = getattr(obj, axis_name)
+        if not _need_slice(slice_obj):
+            return obj
+
+        labels = obj._get_axis(axis)
 
         int_slice = _is_index_slice(slice_obj)
 
@@ -569,23 +570,22 @@ def _get_slice_axis(self, slice_obj, axis=0):
                 raise
 
         if null_slice or position_slice:
-            slicer = slice_obj
+            indexer = slice_obj
         else:
             try:
-                i, j = labels.slice_locs(start, stop)
-                slicer = slice(i, j, slice_obj.step)
+                indexer = labels.slice_indexer(start, stop, slice_obj.step)
             except Exception:
                 if _is_index_slice(slice_obj):
                     if labels.inferred_type == 'integer':
                         raise
-                    slicer = slice_obj
+                    indexer = slice_obj
                 else:
                     raise
 
-        if not _need_slice(slice_obj):
-            return obj
-
-        return self._slice(slicer, axis=axis)
+        if isinstance(indexer, slice):
+            return self._slice(indexer, axis=axis)
+        else:
+            return self.obj.take(indexer, axis=axis)
 
 # 32-bit floating point machine epsilon
 _eps = np.finfo('f4').eps
diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py
@@ -1,6 +1,6 @@
 # pylint: disable-msg=W0612,E1101
 from copy import deepcopy
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, time
 from StringIO import StringIO
 import cPickle as pickle
 import operator
@@ -4552,7 +4552,6 @@ def test_asfreq(self):
         self.assert_(result is not zero_length)
 
     def test_asfreq_datetimeindex(self):
-        from pandas import DatetimeIndex
         df = DataFrame({'A': [1, 2, 3]},
                        index=[datetime(2011, 11, 01), datetime(2011, 11, 2),
                               datetime(2011, 11, 3)])
@@ -4562,6 +4561,44 @@ def test_asfreq_datetimeindex(self):
         ts = df['A'].asfreq('B')
         self.assert_(isinstance(ts.index, DatetimeIndex))
 
+    def test_attime_betweentime_datetimeindex(self):
+        index = pan.date_range("2012-01-01", "2012-01-05", freq='30min')
+        df = DataFrame(randn(len(index), 5), index=index)
+
+        result = df.at_time(time(12, 0, 0))
+        expected = df.ix[time(12, 0, 0)]
+        assert_frame_equal(result, expected)
+        self.assert_(len(result) == 4)
+
+        result = df.between_time(time(13, 0, 0), time(14, 0, 0))
+        expected = df.ix[time(13, 0, 0):time(14, 0, 0)]
+        assert_frame_equal(result, expected)
+        self.assert_(len(result) == 12)
+
+        result = df.copy()
+        result.ix[time(12, 0, 0)] = 0
+        result = result.ix[time(12, 0, 0)]
+        expected = df.ix[time(12, 0, 0)].copy()
+        expected.ix[:] = 0
+        assert_frame_equal(result, expected)
+
+        result = df.copy()
+        result.ix[time(12, 0, 0)] = 0
+        result.ix[time(12, 0, 0)] = df # also checks reindexing
+        assert_frame_equal(result, df)
+
+        result = df.copy()
+        result.ix[time(13, 0, 0):time(14, 0, 0)] = 0
+        result = result.ix[time(13, 0, 0):time(14, 0, 0)]
+        expected = df.ix[time(13, 0, 0):time(14, 0, 0)].copy()
+        expected.ix[:] = 0
+        assert_frame_equal(result, expected)
+
+        result = df.copy()
+        result.ix[time(13, 0, 0):time(14, 0, 0)] = 0
+        result.ix[time(13, 0, 0):time(14, 0, 0)] = df # also checks reindexing
+        assert_frame_equal(result, df)
+
     def test_as_matrix(self):
         frame = self.frame
         mat = frame.as_matrix()
diff --git a/pandas/tseries/index.py b/pandas/tseries/index.py
@@ -1152,6 +1152,20 @@ def _get_string_slice(self, key):
         loc = self._partial_date_slice(reso, parsed)
         return loc
 
+    def slice_indexer(self, start=None, end=None, step=None):
+        """
+        Index.slice_indexer, customized to handle time slicing
+        """
+        if isinstance(start, time) and isinstance(end, time):
+            if step is not None and step != 1:
+                raise ValueError('Must have step size of 1 with time slices')
+            return self.indexer_between_time(start, end)
+
+        if isinstance(start, time) or isinstance(end, time):
+            raise KeyError('Cannot mix time and non-time slice keys')
+
+        return Index.slice_indexer(self, start, end, step)
+
     def slice_locs(self, start=None, end=None):
         """
         Index.slice_locs, customized to handle partial ISO-8601 string slicing
@@ -1172,6 +1186,9 @@ def slice_locs(self, start=None, end=None):
             except KeyError:
                 pass
 
+        if isinstance(start, time) or isinstance(end, time):
+            raise KeyError('Cannot use slice_locs with time slice keys')
+
         return Index.slice_locs(self, start, end)
 
     def __getitem__(self, key):
