@@ -788,6 +788,58 @@ def test_loc_setitem_multiindex(self):
788788 result = df .loc [(t ,n ),'X' ]
789789 self .assertEqual (result ,3 )
790790
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+
791843 def test_loc_setitem_dups (self ):
792844
793845 # GH 6541
0 commit comments