ENH: between_time, at_time accept axis parameter (#21799)

Author: yrhooke

doc/source/whatsnew/v0.24.0.rst

@@ -291,6 +291,7 @@ Other Enhancements
 - :meth:`read_excel()` now accepts ``usecols`` as a list of column names or callable (:issue:`18273`)
 - :meth:`MultiIndex.to_flat_index` has been added to flatten multiple levels into a single-level :class:`Index` object.
 - :meth:`DataFrame.to_stata` and :class:` pandas.io.stata.StataWriter117` can write mixed sting columns to Stata strl format (:issue:`23633`)
+- :meth:`DataFrame.between_time` and :meth:`DataFrame.at_time` have gained the an ``axis`` parameter (:issue: `8839`)
 
 .. _whatsnew_0240.api_breaking:
 

pandas/core/generic.py

@@ -7381,7 +7381,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).
 
@@ -7393,6 +7393,10 @@ def at_time(self, time, asof=False):
         Parameters
         ----------
         time : datetime.time or string
+        axis : {0 or 'index', 1 or 'columns'}, default 0
+
+            .. versionadded:: 0.24.0
+
 
         Returns
         -------
@@ -7422,14 +7426,20 @@ 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)
+
+        index = self._get_axis(axis)
         try:
-            indexer = self.index.indexer_at_time(time, asof=asof)
-            return self._take(indexer)
+            indexer = index.indexer_at_time(time, asof=asof)
         except AttributeError:
             raise TypeError('Index must be DatetimeIndex')
 
+        return self._take(indexer, axis=axis)
+
     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).
 
@@ -7447,6 +7457,9 @@ 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 : {0 or 'index', 1 or 'columns'}, default 0
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -7484,14 +7497,20 @@ 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)
+
+        index = self._get_axis(axis)
         try:
-            indexer = self.index.indexer_between_time(
+            indexer = index.indexer_between_time(
                 start_time, end_time, include_start=include_start,
                 include_end=include_end)
-            return self._take(indexer)
         except AttributeError:
             raise TypeError('Index must be DatetimeIndex')
 
+        return self._take(indexer, axis=axis)
+
     def resample(self, rule, how=None, axis=0, fill_method=None, closed=None,
                  label=None, convention='start', kind=None, loffset=None,
                  limit=None, base=0, on=None, level=None):

pandas/tests/frame/test_timeseries.py

@@ -25,6 +25,11 @@
 from pandas.tests.frame.common import TestData
 
 
+@pytest.fixture(params=product([True, False], [True, False]))
+def close_open_fixture(request):
+    return request.param
+
+
 class TestDataFrameTimeSeriesMethods(TestData):
 
     def test_diff(self):
@@ -638,33 +643,49 @@ def test_at_time_raises(self):
         with pytest.raises(TypeError):  # index is not a DatetimeIndex
             df.at_time('00:00')
 
-    def test_between_time(self):
+    @pytest.mark.parametrize('axis', ['index', 'columns', 0, 1])
+    def test_at_time_axis(self, axis):
+        # issue 8839
+        rng = date_range('1/1/2000', '1/5/2000', freq='5min')
+        ts = DataFrame(np.random.randn(len(rng), len(rng)))
+        ts.index, ts.columns = rng, rng
+
+        indices = rng[(rng.hour == 9) & (rng.minute == 30) & (rng.second == 0)]
+
+        if axis in ['index', 0]:
+            expected = ts.loc[indices, :]
+        elif axis in ['columns', 1]:
+            expected = ts.loc[:, indices]
+
+        result = ts.at_time('9:30', axis=axis)
+        assert_frame_equal(result, expected)
+
+    def test_between_time(self, close_open_fixture):
         rng = date_range('1/1/2000', '1/5/2000', freq='5min')
         ts = DataFrame(np.random.randn(len(rng), 2), index=rng)
         stime = time(0, 0)
         etime = time(1, 0)
-
-        close_open = product([True, False], [True, False])
-        for inc_start, inc_end in close_open:
-            filtered = ts.between_time(stime, etime, inc_start, inc_end)
-            exp_len = 13 * 4 + 1
-            if not inc_start:
-                exp_len -= 5
-            if not inc_end:
-                exp_len -= 4
-
-            assert len(filtered) == exp_len
-            for rs in filtered.index:
-                t = rs.time()
-                if inc_start:
-                    assert t >= stime
-                else:
-                    assert t > stime
-
-                if inc_end:
-                    assert t <= etime
-                else:
-                    assert t < etime
+        inc_start, inc_end = close_open_fixture
+
+        filtered = ts.between_time(stime, etime, inc_start, inc_end)
+        exp_len = 13 * 4 + 1
+        if not inc_start:
+            exp_len -= 5
+        if not inc_end:
+            exp_len -= 4
+
+        assert len(filtered) == exp_len
+        for rs in filtered.index:
+            t = rs.time()
+            if inc_start:
+                assert t >= stime
+            else:
+                assert t > stime
+
+            if inc_end:
+                assert t <= etime
+            else:
+                assert t < etime
 
         result = ts.between_time('00:00', '01:00')
         expected = ts.between_time(stime, etime)
@@ -676,34 +697,66 @@ def test_between_time(self):
         stime = time(22, 0)
         etime = time(9, 0)
 
-        close_open = product([True, False], [True, False])
-        for inc_start, inc_end in close_open:
-            filtered = ts.between_time(stime, etime, inc_start, inc_end)
-            exp_len = (12 * 11 + 1) * 4 + 1
-            if not inc_start:
-                exp_len -= 4
-            if not inc_end:
-                exp_len -= 4
-
-            assert len(filtered) == exp_len
-            for rs in filtered.index:
-                t = rs.time()
-                if inc_start:
-                    assert (t >= stime) or (t <= etime)
-                else:
-                    assert (t > stime) or (t <= etime)
-
-                if inc_end:
-                    assert (t <= etime) or (t >= stime)
-                else:
-                    assert (t < etime) or (t >= stime)
+        filtered = ts.between_time(stime, etime, inc_start, inc_end)
+        exp_len = (12 * 11 + 1) * 4 + 1
+        if not inc_start:
+            exp_len -= 4
+        if not inc_end:
+            exp_len -= 4
+
+        assert len(filtered) == exp_len
+        for rs in filtered.index:
+            t = rs.time()
+            if inc_start:
+                assert (t >= stime) or (t <= etime)
+            else:
+                assert (t > stime) or (t <= etime)
+
+            if inc_end:
+                assert (t <= etime) or (t >= stime)
+            else:
+                assert (t < etime) or (t >= stime)
 
     def test_between_time_raises(self):
         # GH20725
         df = pd.DataFrame([[1, 2, 3], [4, 5, 6]])
         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, axis):
+        # issue 8839
+        rng = date_range('1/1/2000', periods=100, freq='10min')
+        ts = DataFrame(np.random.randn(len(rng), len(rng)))
+        stime, etime = ('08:00:00', '09:00:00')
+        exp_len = 7
+
+        if axis in ['index', 0]:
+            ts.index = rng
+            assert len(ts.between_time(stime, etime)) == exp_len
+            assert len(ts.between_time(stime, etime, axis=0)) == exp_len
+
+        if axis in ['columns', 1]:
+            ts.columns = rng
+            selected = ts.between_time(stime, etime, axis=1).columns
+            assert len(selected) == exp_len
+
+    def test_between_time_axis_raises(self, axis):
+        # issue 8839
+        rng = date_range('1/1/2000', periods=100, freq='10min')
+        mask = np.arange(0, len(rng))
+        rand_data = np.random.randn(len(rng), len(rng))
+        ts = DataFrame(rand_data, index=rng, columns=rng)
+        stime, etime = ('08:00:00', '09:00:00')
+
+        if axis in ['columns', 1]:
+            ts.index = mask
+            pytest.raises(TypeError, ts.between_time, stime, etime)
+            pytest.raises(TypeError, ts.between_time, stime, etime, axis=0)
+
+        if axis in ['index', 0]:
+            ts.columns = mask
+            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,17 @@ 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):
+        # issue 8839
+        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