ENH: between_time, at_time accept axis parameter

Author: yrhooke

pandas/core/generic.py

@@ -6760,7 +6760,7 @@ def asfreq(self, freq, method=None, how=None, normalize=False,
         return asfreq(self, freq, method=method, how=how, normalize=normalize,
                       fill_value=fill_value)
 
-    def at_time(self, time, asof=False):
+    def at_time(self, time, asof=False, axis=None):
         """
         Select values at particular time of day (e.g. 9:30AM).
 
@@ -6772,6 +6772,7 @@ def at_time(self, time, asof=False):
         Parameters
         ----------
         time : datetime.time or string
+        axis : int or string axis name, optional
 
         Returns
         -------
@@ -6801,14 +6802,19 @@ def at_time(self, time, asof=False):
         DatetimeIndex.indexer_at_time : Get just the index locations for
             values at particular time of the day
         """
+        if axis is None:
+            axis = self._stat_axis_number
+        axis = self._get_axis_number(axis)
+
         try:
-            indexer = self.index.indexer_at_time(time, asof=asof)
-            return self._take(indexer)
+            index = self._get_axis(axis)
+            indexer = index.indexer_at_time(time, asof=asof)
+            return self._take(indexer, axis=axis)
         except AttributeError:
             raise TypeError('Index must be DatetimeIndex')
 
     def between_time(self, start_time, end_time, include_start=True,
-                     include_end=True):
+                     include_end=True, axis=None):
         """
         Select values between particular times of the day (e.g., 9:00-9:30 AM).
 
@@ -6826,6 +6832,7 @@ def between_time(self, start_time, end_time, include_start=True,
         end_time : datetime.time or string
         include_start : boolean, default True
         include_end : boolean, default True
+        axis : int or string axis name, optional
 
         Returns
         -------
@@ -6863,11 +6870,16 @@ def between_time(self, start_time, end_time, include_start=True,
         DatetimeIndex.indexer_between_time : Get just the index locations for
             values between particular times of the day
         """
+        if axis is None:
+            axis = self._stat_axis_number
+        axis = self._get_axis_number(axis)
+
         try:
-            indexer = self.index.indexer_between_time(
+            index = self._get_axis(axis)
+            indexer = index.indexer_between_time(
                 start_time, end_time, include_start=include_start,
                 include_end=include_end)
-            return self._take(indexer)
+            return self._take(indexer, axis=axis)
         except AttributeError:
             raise TypeError('Index must be DatetimeIndex')
 

pandas/tests/frame/test_timeseries.py

@@ -640,6 +640,22 @@ def test_at_time_raises(self):
         with pytest.raises(TypeError):  # index is not a DatetimeIndex
             df.at_time('00:00')
 
+    def test_at_time_axis(self):
+        rng = date_range('1/1/2000', '1/5/2000', freq='5min')
+        ts = DataFrame(np.random.randn(len(rng), len(rng)))
+
+        indices = rng[(rng.hour == 9) & (rng.minute == 30) & (rng.second == 0)]
+
+        ts.index = rng
+        expected = ts.loc[indices]
+        result = ts.at_time('9:30', axis=0)
+        assert_frame_equal(result, expected)
+
+        ts.columns = rng
+        expected = ts.loc[:,indices]
+        result = ts.at_time('9:30', axis=1)
+        assert_frame_equal(result, expected)
+
     def test_between_time(self):
         rng = date_range('1/1/2000', '1/5/2000', freq='5min')
         ts = DataFrame(np.random.randn(len(rng), 2), index=rng)
@@ -706,6 +722,40 @@ def test_between_time_raises(self):
         with pytest.raises(TypeError):  # index is not a DatetimeIndex
             df.between_time(start_time='00:00', end_time='12:00')
 
+    def test_between_time_axis(self):
+        rng = date_range('1/1/2000', periods=100, freq='10min')
+        blank = np.arange(0,len(rng))
+        stime, etime = ('08:00:00', '09:00:00')
+        dimn = (len(rng), len(rng))
+        exp_len = 7
+        
+        for time_index, time_col in product([True, False], [True, False]):
+            if time_index:
+                index = rng
+            else:
+                index = blank
+            if time_col:
+                col = rng
+            else:
+                col = blank
+
+            ts = DataFrame(np.random.randn(*dimn), index=index, columns=col)
+
+            if time_index:
+                assert len(ts.between_time(stime, etime)) == exp_len
+                assert len(ts.between_time(stime, etime, axis=0)) == exp_len
+            else:
+                pytest.raises(TypeError, ts.between_time, stime, etime)
+                pytest.raises(TypeError, ts.between_time, stime, etime, 
+                                                                  axis=0)
+
+            if time_col:
+                selected = ts.between_time(stime, etime, axis=1).columns
+                assert len(selected) == exp_len
+            else:
+                pytest.raises(TypeError, ts.between_time, stime, etime, 
+                                                                  axis=1)
+
     def test_operation_on_NaT(self):
         # Both NaT and Timestamp are in DataFrame.
         df = pd.DataFrame({'foo': [pd.NaT, pd.NaT,

pandas/tests/series/test_timeseries.py

@@ -817,6 +817,16 @@ def test_between_time_formats(self):
         for time_string in strings:
             assert len(ts.between_time(*time_string)) == expected_length
 
+    def test_between_time_axis(self):
+        rng = date_range('1/1/2000', periods=100, freq='10min')
+        ts = Series(np.random.randn(len(rng)), index=rng)
+        stime, etime = ('08:00:00', '09:00:00')
+        expected_length = 7
+
+        assert len(ts.between_time(stime, etime)) == expected_length
+        assert len(ts.between_time(stime, etime, axis=0)) == expected_length
+        pytest.raises(ValueError, ts.between_time, stime, etime, axis=1)
+
     def test_to_period(self):
         from pandas.core.indexes.period import period_range