BUG: DatetimeIndex.to_period buggy with tz #2232

Author: changhiskhan

pandas/src/datetime.pxd

@@ -1,4 +1,4 @@
-from numpy cimport int64_t, int32_t, npy_int64, npy_int32
+from numpy cimport int64_t, int32_t, npy_int64, npy_int32, ndarray
 from cpython cimport PyObject
 
 from cpython cimport PyUnicode_Check, PyUnicode_AsASCIIString
@@ -140,3 +140,6 @@ cdef inline int _cstring_to_dts(char *val, int length,
                                      NPY_UNSAFE_CASTING,
                                      dts, &islocal, &out_bestunit, &special)
     return result
+
+cdef ndarray[int64_t] localize_dt64arr_to_period(ndarray[int64_t] stamps,
+                                                 int freq, object tz)

pandas/src/datetime.pyx

@@ -1603,3 +1603,74 @@ cpdef normalize_date(object dt):
         return datetime(dt.year, dt.month, dt.day)
     else:
         raise TypeError('Unrecognized type: %s' % type(dt))
+
+
+cdef extern from "period.h":
+    int64_t get_period_ordinal(int year, int month, int day,
+                          int hour, int minute, int second,
+                          int freq) except INT32_MIN
+
+cdef ndarray[int64_t] localize_dt64arr_to_period(ndarray[int64_t] stamps,
+                                                 int freq, object tz):
+    cdef:
+        Py_ssize_t n = len(stamps)
+        ndarray[int64_t] result = np.empty(n, dtype=np.int64)
+        ndarray[int64_t] trans, deltas, pos
+        pandas_datetimestruct dts
+
+    if not have_pytz:
+        raise Exception('Could not find pytz module')
+
+    if _is_utc(tz):
+        for i in range(n):
+            if stamps[i] == NPY_NAT:
+                result[i] = NPY_NAT
+                continue
+            pandas_datetime_to_datetimestruct(stamps[i], PANDAS_FR_ns, &dts)
+            result[i] = get_period_ordinal(dts.year, dts.month, dts.day,
+                                           dts.hour, dts.min, dts.sec, freq)
+
+    elif _is_tzlocal(tz):
+        for i in range(n):
+            if stamps[i] == NPY_NAT:
+                result[i] = NPY_NAT
+                continue
+            pandas_datetime_to_datetimestruct(stamps[i], PANDAS_FR_ns,
+                                              &dts)
+            dt = datetime(dts.year, dts.month, dts.day, dts.hour,
+                          dts.min, dts.sec, dts.us, tz)
+            delta = int(total_seconds(_get_utcoffset(tz, dt))) * 1000000000
+            pandas_datetime_to_datetimestruct(stamps[i] + delta,
+                                              PANDAS_FR_ns, &dts)
+            result[i] = get_period_ordinal(dts.year, dts.month, dts.day,
+                                           dts.hour, dts.min, dts.sec, freq)
+    else:
+        # Adjust datetime64 timestamp, recompute datetimestruct
+        trans = _get_transitions(tz)
+        deltas = _get_deltas(tz)
+        _pos = trans.searchsorted(stamps, side='right') - 1
+        if _pos.dtype != np.int64:
+            _pos = _pos.astype(np.int64)
+        pos = _pos
+
+        # statictzinfo
+        if not hasattr(tz, '_transition_info'):
+            for i in range(n):
+                if stamps[i] == NPY_NAT:
+                    result[i] = NPY_NAT
+                    continue
+                pandas_datetime_to_datetimestruct(stamps[i] + deltas[0],
+                                                  PANDAS_FR_ns, &dts)
+                result[i] = get_period_ordinal(dts.year, dts.month, dts.day,
+                                               dts.hour, dts.min, dts.sec, freq)
+        else:
+            for i in range(n):
+                if stamps[i] == NPY_NAT:
+                    result[i] = NPY_NAT
+                    continue
+                pandas_datetime_to_datetimestruct(stamps[i] + deltas[pos[i]],
+                                                  PANDAS_FR_ns, &dts)
+                result[i] = get_period_ordinal(dts.year, dts.month, dts.day,
+                                               dts.hour, dts.min, dts.sec, freq)
+
+    return result

pandas/src/plib.pyx

@@ -1,5 +1,4 @@
 # cython: profile=False
-
 cimport numpy as np
 import numpy as np
 
@@ -24,7 +23,6 @@ import_array()
 # import datetime C API
 PyDateTime_IMPORT
 
-
 cdef extern from "period.h":
     ctypedef struct date_info:
         int64_t absdate
@@ -103,7 +101,7 @@ cdef inline int64_t remove_mult(int64_t period_ord_w_mult, int64_t mult):
 
     return period_ord_w_mult * mult + 1;
 
-def dt64arr_to_periodarr(ndarray[int64_t] dtarr, int freq):
+def dt64arr_to_periodarr(ndarray[int64_t] dtarr, int freq, tz=None):
     """
     Convert array of datetime64 values (passed in as 'i8' dtype) to a set of
     periods corresponding to desired frequency, per period convention.
@@ -117,10 +115,13 @@ def dt64arr_to_periodarr(ndarray[int64_t] dtarr, int freq):
 
     out = np.empty(l, dtype='i8')
 
-    for i in range(l):
-        pandas_datetime_to_datetimestruct(dtarr[i], PANDAS_FR_ns, &dts)
-        out[i] = get_period_ordinal(dts.year, dts.month, dts.day,
-                                    dts.hour, dts.min, dts.sec, freq)
+    if tz is None:
+        for i in range(l):
+            pandas_datetime_to_datetimestruct(dtarr[i], PANDAS_FR_ns, &dts)
+            out[i] = get_period_ordinal(dts.year, dts.month, dts.day,
+                                        dts.hour, dts.min, dts.sec, freq)
+    else:
+        out = localize_dt64arr_to_period(dtarr, freq, tz)
     return out
 
 def periodarr_to_dt64arr(ndarray[int64_t] periodarr, int freq):
@@ -353,4 +354,3 @@ cdef accessor _get_accessor_func(int code):
         return &pweekday
     else:
         raise ValueError('Unrecognized code: %s' % code)
-

pandas/tseries/index.py

@@ -703,7 +703,7 @@ def to_period(self, freq=None):
         if freq is None:
             freq = get_period_alias(self.freqstr)
 
-        return PeriodIndex(self.values, freq=freq)
+        return PeriodIndex(self.values, freq=freq, tz=self.tz)
 
     def order(self, return_indexer=False, ascending=True):
         """

pandas/tseries/period.py

@@ -421,12 +421,12 @@ def _get_ordinals(data, freq):
         return lib.map_infer(data, f)
 
 
-def dt64arr_to_periodarr(data, freq):
+def dt64arr_to_periodarr(data, freq, tz):
     if data.dtype != np.dtype('M8[ns]'):
         raise ValueError('Wrong dtype: %s' % data.dtype)
 
     base, mult = _gfc(freq)
-    return plib.dt64arr_to_periodarr(data.view('i8'), base)
+    return plib.dt64arr_to_periodarr(data.view('i8'), base, tz)
 
 # --- Period index sketch
 def _period_index_cmp(opname):
@@ -494,6 +494,8 @@ class PeriodIndex(Int64Index):
     hour : int or array, default None
     minute : int or array, default None
     second : int or array, default None
+    tz : object, default None
+        Timezone for converting datetime64 data to Periods
 
     Examples
     --------
@@ -514,7 +516,8 @@ def __new__(cls, data=None, ordinal=None,
                 freq=None, start=None, end=None, periods=None,
                 copy=False, name=None,
                 year=None, month=None, quarter=None, day=None,
-                hour=None, minute=None, second=None):
+                hour=None, minute=None, second=None,
+                tz=None):
 
         freq = _freq_mod.get_standard_freq(freq)
 
@@ -531,9 +534,9 @@ def __new__(cls, data=None, ordinal=None,
             else:
                 fields = [year, month, quarter, day, hour, minute, second]
                 data, freq = cls._generate_range(start, end, periods,
-                                                    freq, fields)
+                                                 freq, fields)
         else:
-            ordinal, freq = cls._from_arraylike(data, freq)
+            ordinal, freq = cls._from_arraylike(data, freq, tz)
             data = np.array(ordinal, dtype=np.int64, copy=False)
 
         subarr = data.view(cls)
@@ -562,7 +565,7 @@ def _generate_range(cls, start, end, periods, freq, fields):
         return subarr, freq
 
     @classmethod
-    def _from_arraylike(cls, data, freq):
+    def _from_arraylike(cls, data, freq, tz):
         if not isinstance(data, np.ndarray):
             if np.isscalar(data) or isinstance(data, Period):
                 raise ValueError('PeriodIndex() must be called with a '
@@ -608,7 +611,7 @@ def _from_arraylike(cls, data, freq):
                                       'inferred from first element'))
 
                 if np.issubdtype(data.dtype, np.datetime64):
-                    data = dt64arr_to_periodarr(data, freq)
+                    data = dt64arr_to_periodarr(data, freq, tz)
                 elif data.dtype == np.int64:
                     pass
                 else:
@@ -1219,4 +1222,3 @@ def period_range(start=None, end=None, periods=None, freq='D', name=None):
     """
     return PeriodIndex(start=start, end=end, periods=periods,
                        freq=freq, name=name)
-

setup.py

@@ -563,6 +563,7 @@ def srcpath(name=None, suffix='.pyx', subdir='src'):
     plib_depends = [srcpath(f, suffix='.pyx')
                     for f in plib_depends]
     plib_depends.append('pandas/src/util.pxd')
+    plib_depends.append('pandas/src/datetime.pxd')
 else:
     tseries_depends = []
     plib_depends = []