BUG: Slicing and setitem inconsistent with getitem using DateTimeIndex with timezone

Author: stephenwlin

pandas/tests/test_series.py

@@ -2544,13 +2544,75 @@ def test_asof(self):
  d = self.ts.index[0] - datetools.bday
  self.assert_(np.isnan(self.ts.asof(d)))
 
+ def test_getitem_setitem_datetimeindex(self):
+ from pandas import date_range
+ N = 50
+ # testing with timezone, GH #2785
+ rng = date_range('1/1/1990', periods=N, freq='H', tz='US/Eastern')
+ ts = Series(np.random.randn(N), index=rng)
+
+ result = ts["1990-01-01 04:00:00"]
+ expected = ts[4]
+ self.assert_(result == expected)
+
+ result = ts.copy()
+ result["1990-01-01 04:00:00"] = 0
+ result["1990-01-01 04:00:00"] = ts[4]
+ assert_series_equal(result, ts)
+
+ result = ts["1990-01-01 04:00:00":"1990-01-01 07:00:00"]
+ expected = ts[4:8]
+ assert_series_equal(result, expected)
+
+ result = ts.copy()
+ result["1990-01-01 04:00:00":"1990-01-01 07:00:00"] = 0
+ result["1990-01-01 04:00:00":"1990-01-01 07:00:00"] = ts[4:8]
+ assert_series_equal(result, ts)
+
+ lb = "1990-01-01 04:00:00"
+ rb = "1990-01-01 07:00:00"
+ result = ts[(ts.index >= lb) & (ts.index <= rb)]
+ expected = ts[4:8]
+ assert_series_equal(result, expected)
+
+ result = ts[ts.index[4]]
+ expected = ts[4]
+ self.assert_(result == expected)
+
+ result = ts[ts.index[4:8]]
+ expected = ts[4:8]
+ assert_series_equal(result, expected)
+
+ result = ts.copy()
+ result[ts.index[4:8]] = 0
+ result[4:8] = ts[4:8]
+ assert_series_equal(result, ts)
+
+ # also test partial date slicing
+ result = ts["1990-01-02"]
+ expected = ts[24:48]
+ assert_series_equal(result, expected)
+
+ result = ts.copy()
+ result["1990-01-02"] = 0
+ result["1990-01-02"] = ts[24:48]
+ assert_series_equal(result, ts)
+
  def test_getitem_setitem_periodindex(self):
  from pandas import period_range, Period
- # array or list or dates
  N = 50
  rng = period_range('1/1/1990', periods=N, freq='H')
  ts = Series(np.random.randn(N), index=rng)
 
+ result = ts["1990-01-01 04"]
+ expected = ts[4]
+ self.assert_(result == expected)
+
+ result = ts.copy()
+ result["1990-01-01 04"] = 0
+ result["1990-01-01 04"] = ts[4]
+ assert_series_equal(result, ts)
+
  result = ts["1990-01-01 04":"1990-01-01 07"]
  expected = ts[4:8]
  assert_series_equal(result, expected)
@@ -2567,6 +2629,10 @@ def test_getitem_setitem_periodindex(self):
  assert_series_equal(result, expected)
 
  # GH 2782
+ result = ts[ts.index[4]]
+ expected = ts[4]
+ self.assert_(result == expected)
+
  result = ts[ts.index[4:8]]
  expected = ts[4:8]
  assert_series_equal(result, expected)

pandas/tseries/index.py

@@ -71,6 +71,8 @@ def wrapper(self, other):
  other = _to_m8(other)
  elif isinstance(other, list):
  other = DatetimeIndex(other)
+ elif isinstance(other, basestring):
+ other = _to_m8(Timestamp(other, tz=self.tz))
  elif not isinstance(other, np.ndarray):
  other = _ensure_datetime64(other)
  result = func(other)
@@ -1042,34 +1044,36 @@ def _partial_date_slice(self, reso, parsed):
  'time series.')
 
  if reso == 'year':
- t1 = Timestamp(datetime(parsed.year, 1, 1))
- t2 = Timestamp(datetime(parsed.year, 12, 31))
+ t1 = Timestamp(datetime(parsed.year, 1, 1), tz=self.tz)
+ t2 = Timestamp(datetime(parsed.year, 12, 31), tz=self.tz)
  elif reso == 'month':
  d = tslib.monthrange(parsed.year, parsed.month)[1]
- t1 = Timestamp(datetime(parsed.year, parsed.month, 1))
- t2 = Timestamp(datetime(parsed.year, parsed.month, d))
+ t1 = Timestamp(datetime(parsed.year, parsed.month, 1), tz=self.tz)
+ t2 = Timestamp(datetime(parsed.year, parsed.month, d), tz=self.tz)
  elif reso == 'quarter':
  qe = (((parsed.month - 1) + 2) % 12) + 1 # two months ahead
  d = tslib.monthrange(parsed.year, qe)[1] # at end of month
- t1 = Timestamp(datetime(parsed.year, parsed.month, 1))
- t2 = Timestamp(datetime(parsed.year, qe, d))
+ t1 = Timestamp(datetime(parsed.year, parsed.month, 1), tz=self.tz)
+ t2 = Timestamp(datetime(parsed.year, qe, d), tz=self.tz)
  elif reso == 'day' and self._resolution < Resolution.RESO_DAY:
  st = datetime(parsed.year, parsed.month, parsed.day)
- t1 = Timestamp(st)
+ t1 = Timestamp(st, tz=self.tz)
  t2 = st + offsets.Day()
- t2 = Timestamp(Timestamp(t2).value - 1)
+ t2 = Timestamp(Timestamp(t2, tz=self.tz).value - 1)
  elif (reso == 'hour' and
  self._resolution < Resolution.RESO_HR):
  st = datetime(parsed.year, parsed.month, parsed.day,
  hour=parsed.hour)
- t1 = Timestamp(st)
- t2 = Timestamp(Timestamp(st + offsets.Hour()).value - 1)
+ t1 = Timestamp(st, tz=self.tz)
+ t2 = Timestamp(Timestamp(st + offsets.Hour(),
+ tz=self.tz).value - 1)
  elif (reso == 'minute' and
  self._resolution < Resolution.RESO_MIN):
  st = datetime(parsed.year, parsed.month, parsed.day,
  hour=parsed.hour, minute=parsed.minute)
- t1 = Timestamp(st)
- t2 = Timestamp(Timestamp(st + offsets.Minute()).value - 1)
+ t1 = Timestamp(st, tz=self.tz)
+ t2 = Timestamp(Timestamp(st + offsets.Minute(),
+ tz=self.tz).value - 1)
  else:
  raise KeyError
 
@@ -1091,7 +1095,6 @@ def get_value(self, series, key):
  try:
  return Index.get_value(self, series, key)
  except KeyError:
-
  try:
  loc = self._get_string_slice(key)
  return series[loc]
@@ -1102,11 +1105,11 @@ def get_value(self, series, key):
  locs = self.indexer_at_time(key)
  return series.take(locs)
 
- if isinstance(key, basestring):
- stamp = Timestamp(key, tz=self.tz)
- else:
- stamp = Timestamp(key)
  try:
+ if isinstance(key, basestring):
+ stamp = Timestamp(key, tz=self.tz)
+ else:
+ stamp = Timestamp(key)
  return self._engine.get_value(series, stamp)
  except KeyError:
  raise KeyError(stamp)
@@ -1131,15 +1134,18 @@ def get_loc(self, key):
  return self.indexer_at_time(key)
 
  try:
- return self._engine.get_loc(Timestamp(key))
+ if isinstance(key, basestring):
+ stamp = Timestamp(key, tz=self.tz)
+ else:
+ stamp = Timestamp(key)
+ return self._engine.get_loc(stamp)
  except (KeyError, ValueError):
  raise KeyError(key)
 
  def _get_string_slice(self, key):
  freq = getattr(self, 'freqstr',
  getattr(self, 'inferred_freq', None))
- asdt, parsed, reso = parse_time_string(key, freq)
- key = asdt
+ _, parsed, reso = parse_time_string(key, freq)
  loc = self._partial_date_slice(reso, parsed)
  return loc
 
@@ -1617,7 +1623,7 @@ def _to_m8(key):
  '''
  Timestamp-like => dt64
  '''
- if not isinstance(key, datetime):
+ if not isinstance(key, (Timestamp, datetime)):
  # this also converts strings
  key = Timestamp(key)