import numpy as np import pandas as pd from pandas import date_range, Index, DataFrame, Series, Timestamp from pandas.util import testing as tm class TestDatetimeIndex(object): def test_setitem_with_datetime_tz(self): # 16889 # support .loc with alignment and tz-aware DatetimeIndex mask = np.array([True, False, True, False]) idx = date_range('20010101', periods=4, tz='UTC') df = DataFrame({'a': np.arange(4)}, index=idx).astype('float64') result = df.copy() result.loc[mask, :] = df.loc[mask, :] tm.assert_frame_equal(result, df) result = df.copy() result.loc[mask] = df.loc[mask] tm.assert_frame_equal(result, df) idx = date_range('20010101', periods=4) df = DataFrame({'a': np.arange(4)}, index=idx).astype('float64') result = df.copy() result.loc[mask, :] = df.loc[mask, :] tm.assert_frame_equal(result, df) result = df.copy() result.loc[mask] = df.loc[mask] tm.assert_frame_equal(result, df) def test_indexing_with_datetime_tz(self): # 8260 # support datetime64 with tz idx = Index(date_range('20130101', periods=3, tz='US/Eastern'), name='foo') dr = date_range('20130110', periods=3) df = DataFrame({'A': idx, 'B': dr}) df['C'] = idx df.iloc[1, 1] = pd.NaT df.iloc[1, 2] = pd.NaT # indexing result = df.iloc[1] expected = Series([Timestamp('2013-01-02 00:00:00-0500', tz='US/Eastern'), np.nan, np.nan], index=list('ABC'), dtype='object', name=1) tm.assert_series_equal(result, expected) result = df.loc[1] expected = Series([Timestamp('2013-01-02 00:00:00-0500', tz='US/Eastern'), np.nan, np.nan], index=list('ABC'), dtype='object', name=1) tm.assert_series_equal(result, expected) # indexing - fast_xs df = DataFrame({'a': date_range('2014-01-01', periods=10, tz='UTC')}) result = df.iloc[5] expected = Timestamp('2014-01-06 00:00:00+0000', tz='UTC', freq='D') assert result == expected result = df.loc[5] assert result == expected # indexing - boolean result = df[df.a > df.a[3]] expected = df.iloc[4:] tm.assert_frame_equal(result, expected) # indexing - setting an element df = DataFrame(data=pd.to_datetime( ['2015-03-30 20:12:32', '2015-03-12 00:11:11']), columns=['time']) df['new_col'] = ['new', 'old'] df.time = df.set_index('time').index.tz_localize('UTC') v = df[df.new_col == 'new'].set_index('time').index.tz_convert( 'US/Pacific') # trying to set a single element on a part of a different timezone # this converts to object df2 = df.copy() df2.loc[df2.new_col == 'new', 'time'] = v expected = Series([v[0], df.loc[1, 'time']], name='time') tm.assert_series_equal(df2.time, expected) v = df.loc[df.new_col == 'new', 'time'] + pd.Timedelta('1s') df.loc[df.new_col == 'new', 'time'] = v tm.assert_series_equal(df.loc[df.new_col == 'new', 'time'], v) def test_consistency_with_tz_aware_scalar(self): # xef gh-12938 # various ways of indexing the same tz-aware scalar df = Series([Timestamp('2016-03-30 14:35:25', tz='Europe/Brussels')]).to_frame() df = pd.concat([df, df]).reset_index(drop=True) expected = Timestamp('2016-03-30 14:35:25+0200', tz='Europe/Brussels') result = df[0][0] assert result == expected result = df.iloc[0, 0] assert result == expected result = df.loc[0, 0] assert result == expected result = df.iat[0, 0] assert result == expected result = df.at[0, 0] assert result == expected result = df[0].loc[0] assert result == expected result = df[0].at[0] assert result == expected def test_indexing_with_datetimeindex_tz(self): # GH 12050 # indexing on a series with a datetimeindex with tz index = date_range('2015-01-01', periods=2, tz='utc') ser = Series(range(2), index=index, dtype='int64') # list-like indexing for sel in (index, list(index)): # getitem tm.assert_series_equal(ser[sel], ser) # setitem result = ser.copy() result[sel] = 1 expected = Series(1, index=index) tm.assert_series_equal(result, expected) # .loc getitem tm.assert_series_equal(ser.loc[sel], ser) # .loc setitem result = ser.copy() result.loc[sel] = 1 expected = Series(1, index=index) tm.assert_series_equal(result, expected) # single element indexing # getitem assert ser[index[1]] == 1 # setitem result = ser.copy() result[index[1]] = 5 expected = Series([0, 5], index=index) tm.assert_series_equal(result, expected) # .loc getitem assert ser.loc[index[1]] == 1 # .loc setitem result = ser.copy() result.loc[index[1]] = 5 expected = Series([0, 5], index=index) tm.assert_series_equal(result, expected) def test_partial_setting_with_datetimelike_dtype(self): # GH9478 # a datetimeindex alignment issue with partial setting df = DataFrame(np.arange(6.).reshape(3, 2), columns=list('AB'), index=date_range('1/1/2000', periods=3, freq='1H')) expected = df.copy() expected['C'] = [expected.index[0]] + [pd.NaT, pd.NaT] mask = df.A < 1 df.loc[mask, 'C'] = df.loc[mask].index tm.assert_frame_equal(df, expected) def test_loc_setitem_datetime(self): # GH 9516 dt1 = Timestamp('20130101 09:00:00') dt2 = Timestamp('20130101 10:00:00') for conv in [lambda x: x, lambda x: x.to_datetime64(), lambda x: x.to_pydatetime(), lambda x: np.datetime64(x)]: df = DataFrame() df.loc[conv(dt1), 'one'] = 100 df.loc[conv(dt2), 'one'] = 200 expected = DataFrame({'one': [100.0, 200.0]}, index=[dt1, dt2]) tm.assert_frame_equal(df, expected) def test_series_partial_set_datetime(self): # GH 11497 idx = date_range('2011-01-01', '2011-01-02', freq='D', name='idx') ser = Series([0.1, 0.2], index=idx, name='s') result = ser.loc[[Timestamp('2011-01-01'), Timestamp('2011-01-02')]] exp = Series([0.1, 0.2], index=idx, name='s') tm.assert_series_equal(result, exp, check_index_type=True) keys = [Timestamp('2011-01-02'), Timestamp('2011-01-02'), Timestamp('2011-01-01')] exp = Series([0.2, 0.2, 0.1], index=pd.DatetimeIndex(keys, name='idx'), name='s') tm.assert_series_equal(ser.loc[keys], exp, check_index_type=True) keys = [Timestamp('2011-01-03'), Timestamp('2011-01-02'), Timestamp('2011-01-03')] exp = Series([np.nan, 0.2, np.nan], index=pd.DatetimeIndex(keys, name='idx'), name='s') with tm.assert_produces_warning(FutureWarning, check_stacklevel=False): tm.assert_series_equal(ser.loc[keys], exp, check_index_type=True) def test_series_partial_set_period(self): # GH 11497 idx = pd.period_range('2011-01-01', '2011-01-02', freq='D', name='idx') ser = Series([0.1, 0.2], index=idx, name='s') result = ser.loc[[pd.Period('2011-01-01', freq='D'), pd.Period('2011-01-02', freq='D')]] exp = Series([0.1, 0.2], index=idx, name='s') tm.assert_series_equal(result, exp, check_index_type=True) keys = [pd.Period('2011-01-02', freq='D'), pd.Period('2011-01-02', freq='D'), pd.Period('2011-01-01', freq='D')] exp = Series([0.2, 0.2, 0.1], index=pd.PeriodIndex(keys, name='idx'), name='s') tm.assert_series_equal(ser.loc[keys], exp, check_index_type=True) keys = [pd.Period('2011-01-03', freq='D'), pd.Period('2011-01-02', freq='D'), pd.Period('2011-01-03', freq='D')] exp = Series([np.nan, 0.2, np.nan], index=pd.PeriodIndex(keys, name='idx'), name='s') with tm.assert_produces_warning(FutureWarning, check_stacklevel=False): result = ser.loc[keys] tm.assert_series_equal(result, exp)