@@ -640,7 +640,7 @@ def test_at_time_raises(self):
640
640
with pytest .raises (TypeError ): # index is not a DatetimeIndex
641
641
df .at_time ('00:00' )
642
642
643
- @pytest .mark .parametrize ('axis' , ['index' , 'columns' ])
643
+ @pytest .mark .parametrize ('axis' , ['index' , 'columns' , 0 , 1 ])
644
644
def test_at_time_axis (self , axis ):
645
645
# issue 8839
646
646
rng = date_range ('1/1/2000' , '1/5/2000' , freq = '5min' )
@@ -649,9 +649,9 @@ def test_at_time_axis(self, axis):
649
649
650
650
indices = rng [(rng .hour == 9 ) & (rng .minute == 30 ) & (rng .second == 0 )]
651
651
652
- if axis == 'index' :
652
+ if axis in [ 'index' , 0 ] :
653
653
expected = ts .loc [indices , :]
654
- elif axis == 'columns' :
654
+ elif axis in [ 'columns' , 1 ] :
655
655
expected = ts .loc [:, indices ]
656
656
657
657
result = ts .at_time ('9:30' , axis = axis )
@@ -724,20 +724,20 @@ def test_between_time_raises(self):
724
724
df .between_time (start_time = '00:00' , end_time = '12:00' )
725
725
726
726
@pytest .mark .parametrize ('axis' , [
727
- (), 'index' , 'columns' , ('index' , 'columns' )])
727
+ (), 'index' , 'columns' , ('index' , 'columns' ), 0 , 1 , ( 0 , 1 ) ])
728
728
def test_between_time_axis (self , axis ):
729
729
# issue 8839
730
730
rng = date_range ('1/1/2000' , periods = 100 , freq = '10min' )
731
731
ts = DataFrame (np .random .randn (len (rng ), len (rng )))
732
732
stime , etime = ('08:00:00' , '09:00:00' )
733
733
exp_len = 7
734
734
735
- if 'index' in axis :
735
+ if 'index' in axis or 0 in axis :
736
736
ts .index = rng
737
737
assert len (ts .between_time (stime , etime )) == exp_len
738
738
assert len (ts .between_time (stime , etime , axis = 0 )) == exp_len
739
739
740
- if 'columns' in axis :
740
+ if 'columns' in axis or 1 in axis :
741
741
ts .columns = rng
742
742
selected = ts .between_time (stime , etime , axis = 1 ).columns
743
743
assert len (selected ) == exp_len
0 commit comments