BUG: naive datetime keys not localized to DateTimeIndex tz

Author: stephenwlin

pandas/tests/test_series.py

@@ -2576,6 +2576,31 @@ def test_getitem_setitem_datetimeindex(self):
         expected = ts[4:8]
         assert_series_equal(result, expected)
 
+        # repeat all the above with naive datetimes
+        result = ts[datetime(1990, 1, 1, 4)]
+        expected = ts[4]
+        self.assert_(result == expected)
+
+        result = ts.copy()
+        result[datetime(1990, 1, 1, 4)] = 0
+        result[datetime(1990, 1, 1, 4)] = ts[4]
+        assert_series_equal(result, ts)
+
+        result = ts[datetime(1990, 1, 1, 4):datetime(1990, 1, 1, 7)]
+        expected = ts[4:8]
+        assert_series_equal(result, expected)
+
+        result = ts.copy()
+        result[datetime(1990, 1, 1, 4):datetime(1990, 1, 1, 7)] = 0
+        result[datetime(1990, 1, 1, 4):datetime(1990, 1, 1, 7)] = ts[4:8]
+        assert_series_equal(result, ts)
+
+        lb = datetime(1990, 1, 1, 4)
+        rb = datetime(1990, 1, 1, 7)
+        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)
@@ -2610,14 +2635,15 @@ def test_getitem_setitem_datetimeindex(self):
         result["1990-01-01 03:00:00-06:00"] = ts[4]
         assert_series_equal(result, ts)
 
+        # repeat with datetimes
         result = ts.copy()
-        result[datetime(1990, 1, 1, 9, 0, 0, tzinfo=tz('UTC'))] = 0
-        result[datetime(1990, 1, 1, 9, 0, 0, tzinfo=tz('UTC'))] = ts[4]
+        result[datetime(1990, 1, 1, 9, tzinfo=tz('UTC'))] = 0
+        result[datetime(1990, 1, 1, 9, tzinfo=tz('UTC'))] = ts[4]
         assert_series_equal(result, ts)
 
         result = ts.copy()
-        result[datetime(1990, 1, 1, 3, 0, 0, tzinfo=tz('US/Central'))] = 0
-        result[datetime(1990, 1, 1, 3, 0, 0, tzinfo=tz('US/Central'))] = ts[4]
+        result[datetime(1990, 1, 1, 3, tzinfo=tz('US/Central'))] = 0
+        result[datetime(1990, 1, 1, 3, tzinfo=tz('US/Central'))] = ts[4]
         assert_series_equal(result, ts)
 
     def test_getitem_setitem_periodindex(self):

pandas/tseries/index.py

@@ -67,12 +67,11 @@ def _dt_index_cmp(opname):
     def wrapper(self, other):
         func = getattr(super(DatetimeIndex, self), opname)
         if isinstance(other, datetime):
-            func = getattr(self, opname)
-            other = _to_m8(other)
+            other = _to_m8(other, tz=self.tz)
         elif isinstance(other, list):
             other = DatetimeIndex(other)
         elif isinstance(other, basestring):
-            other = _to_m8(Timestamp(other, tz=self.tz))
+            other = _to_m8(other, tz=self.tz)
         elif not isinstance(other, np.ndarray):
             other = _ensure_datetime64(other)
         result = func(other)
@@ -1092,6 +1091,11 @@ def get_value(self, series, key):
         Fast lookup of value from 1-dimensional ndarray. Only use this if you
         know what you're doing
         """
+        if isinstance(key, datetime):
+            # needed to localize naive datetimes
+            stamp = Timestamp(key, tz=self.tz)
+            return self._engine.get_value(series, stamp)
+
         try:
             return Index.get_value(self, series, key)
         except KeyError:
@@ -1106,13 +1110,10 @@ def get_value(self, series, key):
                 return series.take(locs)
 
             try:
-                if isinstance(key, basestring):
-                    stamp = Timestamp(key, tz=self.tz)
-                else:
-                    stamp = Timestamp(key)
+                stamp = Timestamp(key, tz=self.tz)
                 return self._engine.get_value(series, stamp)
-            except KeyError:
-                raise KeyError(stamp)
+            except (KeyError, ValueError):
+                raise KeyError(key)
 
     def get_loc(self, key):
         """
@@ -1122,22 +1123,24 @@ def get_loc(self, key):
         -------
         loc : int
         """
+        if isinstance(key, datetime):
+            # needed to localize naive datetimes
+            stamp = Timestamp(key, tz=self.tz)
+            return self._engine.get_loc(stamp)
+
         try:
-            return self._engine.get_loc(key)
-        except KeyError:
+            return Index.get_loc(self, key)
+        except (KeyError, ValueError):
             try:
                 return self._get_string_slice(key)
             except (TypeError, KeyError, ValueError):
                 pass
 
             if isinstance(key, time):
                 return self.indexer_at_time(key)
-
+            
             try:
-                if isinstance(key, basestring):
-                    stamp = Timestamp(key, tz=self.tz)
-                else:
-                    stamp = Timestamp(key)
+                stamp = Timestamp(key, tz=self.tz)
                 return self._engine.get_loc(stamp)
             except (KeyError, ValueError):
                 raise KeyError(key)
@@ -1256,7 +1259,7 @@ def searchsorted(self, key, side='left'):
         if isinstance(key, np.ndarray):
             key = np.array(key, dtype=_NS_DTYPE, copy=False)
         else:
-            key = _to_m8(key)
+            key = _to_m8(key, tz=self.tz)
 
         return self.values.searchsorted(key, side=side)
 
@@ -1345,7 +1348,7 @@ def insert(self, loc, item):
         new_index : Index
         """
         if isinstance(item, datetime):
-            item = _to_m8(item)
+            item = _to_m8(item, tz=self.tz)
 
         new_index = np.concatenate((self[:loc].asi8,
                                     [item.view(np.int64)],
@@ -1619,13 +1622,13 @@ def bdate_range(start=None, end=None, periods=None, freq='B', tz=None,
                          freq=freq, tz=tz, normalize=normalize, name=name)
 
 
-def _to_m8(key):
+def _to_m8(key, tz=None):
     '''
     Timestamp-like => dt64
     '''
-    if not isinstance(key, (Timestamp, datetime)):
+    if not isinstance(key, Timestamp):
         # this also converts strings
-        key = Timestamp(key)
+        key = Timestamp(key, tz=tz)
 
     return np.int64(tslib.pydt_to_i8(key)).view(_NS_DTYPE)