@@ -788,6 +788,58 @@ def test_loc_setitem_multiindex(self):
788
788
result = df .loc [(t ,n ),'X' ]
789
789
self .assertEqual (result ,3 )
790
790
791
+ def test_indexing_with_datetime_tz (self ):
792
+
793
+ # 8260
794
+ # support datetime64 with tz
795
+
796
+ idx = Index (date_range ('20130101' ,periods = 3 ,tz = 'US/Eastern' ),
797
+ name = 'foo' )
798
+ dr = date_range ('20130110' ,periods = 3 )
799
+ df = DataFrame ({'A' : idx , 'B' : dr })
800
+ df ['C' ] = idx
801
+ df .iloc [1 ,1 ] = pd .NaT
802
+ df .iloc [1 ,2 ] = pd .NaT
803
+
804
+ # indexing
805
+ result = df .iloc [1 ]
806
+ expected = Series ([Timestamp ('2013-01-02 00:00:00-0500' , tz = 'US/Eastern' ), np .nan , np .nan ],
807
+ index = list ('ABC' ), dtype = 'object' , name = 1 )
808
+ assert_series_equal (result , expected )
809
+ result = df .loc [1 ]
810
+ expected = Series ([Timestamp ('2013-01-02 00:00:00-0500' , tz = 'US/Eastern' ), np .nan , np .nan ],
811
+ index = list ('ABC' ), dtype = 'object' , name = 1 )
812
+ assert_series_equal (result , expected )
813
+
814
+ # indexing - fast_xs
815
+ df = DataFrame ({'a' : date_range ('2014-01-01' , periods = 10 , tz = 'UTC' )})
816
+ result = df .iloc [5 ]
817
+ expected = Timestamp ('2014-01-06 00:00:00+0000' , tz = 'UTC' , offset = 'D' )
818
+ self .assertEqual (result , expected )
819
+
820
+ result = df .loc [5 ]
821
+ self .assertEqual (result , expected )
822
+
823
+ # indexing - boolean
824
+ result = df [df .a > df .a [3 ]]
825
+ expected = df .iloc [4 :]
826
+ assert_frame_equal (result , expected )
827
+
828
+ # indexing - setting an element
829
+ df = DataFrame ( data = pd .to_datetime (['2015-03-30 20:12:32' ,'2015-03-12 00:11:11' ]) ,columns = ['time' ] )
830
+ df ['new_col' ]= ['new' ,'old' ]
831
+ df .time = df .set_index ('time' ).index .tz_localize ('UTC' )
832
+ v = df [df .new_col == 'new' ].set_index ('time' ).index .tz_convert ('US/Pacific' )
833
+
834
+ # trying to set a single element on a part of a different timezone
835
+ def f ():
836
+ df .loc [df .new_col == 'new' ,'time' ] = v
837
+ self .assertRaises (ValueError , f )
838
+
839
+ v = df .loc [df .new_col == 'new' ,'time' ] + pd .Timedelta ('1s' )
840
+ df .loc [df .new_col == 'new' ,'time' ] = v
841
+ assert_series_equal (df .loc [df .new_col == 'new' ,'time' ],v )
842
+
791
843
def test_loc_setitem_dups (self ):
792
844
793
845
# GH 6541
0 commit comments