255 lines
8.8 KiB
Python
255 lines
8.8 KiB
Python
|
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)
|