Have DTA._simple_new take dtype instead of tz (#24665)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -296,12 +296,11 @@ def _round(self, freq, mode, ambiguous, nonexistent):
         result = round_nsint64(values, mode, freq)
         result = self._maybe_mask_results(result, fill_value=NaT)
 
-        attribs = self._get_attributes_dict()
-        attribs['freq'] = None
-        if 'tz' in attribs:
-            attribs['tz'] = None
+        dtype = self.dtype
+        if is_datetime64tz_dtype(self):
+            dtype = None
         return self._ensure_localized(
-            self._simple_new(result, **attribs), ambiguous, nonexistent
+            self._simple_new(result, dtype=dtype), ambiguous, nonexistent
         )
 
     @Appender((_round_doc + _round_example).format(op="round"))
@@ -434,8 +433,6 @@ def __getitem__(self, key):
             else:
                 key = lib.maybe_booleans_to_slice(key.view(np.uint8))
 
-        attribs = self._get_attributes_dict()
-
         is_period = is_period_dtype(self)
         if is_period:
             freq = self.freq
@@ -451,17 +448,15 @@ def __getitem__(self, key):
                 #  should preserve `freq` attribute
                 freq = self.freq
 
-        attribs['freq'] = freq
-
         result = getitem(key)
         if result.ndim > 1:
             # To support MPL which performs slicing with 2 dim
             # even though it only has 1 dim by definition
             if is_period:
-                return self._simple_new(result, **attribs)
+                return self._simple_new(result, dtype=self.dtype, freq=freq)
             return result
 
-        return self._simple_new(result, **attribs)
+        return self._simple_new(result, dtype=self.dtype, freq=freq)
 
     def __setitem__(
             self,

pandas/core/arrays/datetimes.py

@@ -35,6 +35,24 @@
 _midnight = time(0, 0)
 
 
+def tz_to_dtype(tz):
+    """
+    Return a datetime64[ns] dtype appropriate for the given timezone.
+
+    Parameters
+    ----------
+    tz : tzinfo or None
+
+    Returns
+    -------
+    np.dtype or Datetime64TZDType
+    """
+    if tz is None:
+        return _NS_DTYPE
+    else:
+        return DatetimeTZDtype(tz=tz)
+
+
 def _to_M8(key, tz=None):
     """
     Timestamp-like => dt64
@@ -305,13 +323,7 @@ def __init__(self, values, dtype=_NS_DTYPE, freq=None, copy=False):
         self._freq = freq
 
     @classmethod
-    def _simple_new(cls, values, freq=None, tz=None):
-        """
-        we require the we have a dtype compat for the values
-        if we are passed a non-dtype compat, then coerce using the constructor
-        """
-        dtype = DatetimeTZDtype(tz=tz) if tz else _NS_DTYPE
-
+    def _simple_new(cls, values, freq=None, dtype=None):
         return cls(values, freq=freq, dtype=dtype)
 
     @classmethod
@@ -328,7 +340,8 @@ def _from_sequence(cls, data, dtype=None, copy=False,
         freq, freq_infer = dtl.validate_inferred_freq(freq, inferred_freq,
                                                       freq_infer)
 
-        result = cls._simple_new(subarr, freq=freq, tz=tz)
+        dtype = tz_to_dtype(tz)
+        result = cls._simple_new(subarr, freq=freq, dtype=dtype)
 
         if inferred_freq is None and freq is not None:
             # this condition precludes `freq_infer`
@@ -395,7 +408,7 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
                     end = end.tz_localize(None)
             # TODO: consider re-implementing _cached_range; GH#17914
             values, _tz = generate_regular_range(start, end, periods, freq)
-            index = cls._simple_new(values, freq=freq, tz=_tz)
+            index = cls._simple_new(values, freq=freq, dtype=tz_to_dtype(_tz))
 
             if tz is not None and index.tz is None:
                 arr = conversion.tz_localize_to_utc(
@@ -418,16 +431,18 @@ def _generate_range(cls, start, end, periods, freq, tz=None,
             arr = np.linspace(
                 0, end.value - start.value,
                 periods, dtype='int64') + start.value
+            dtype = tz_to_dtype(tz)
             index = cls._simple_new(
-                arr.astype('M8[ns]', copy=False), freq=None, tz=tz
+                arr.astype('M8[ns]', copy=False), freq=None, dtype=dtype
             )
 
         if not left_closed and len(index) and index[0] == start:
             index = index[1:]
         if not right_closed and len(index) and index[-1] == end:
             index = index[:-1]
 
-        return cls._simple_new(index.asi8, freq=freq, tz=tz)
+        dtype = tz_to_dtype(tz)
+        return cls._simple_new(index.asi8, freq=freq, dtype=dtype)
 
     # -----------------------------------------------------------------
     # DatetimeLike Interface
@@ -806,7 +821,8 @@ def tz_convert(self, tz):
                             'tz_localize to localize')
 
         # No conversion since timestamps are all UTC to begin with
-        return self._simple_new(self.asi8, tz=tz, freq=self.freq)
+        dtype = tz_to_dtype(tz)
+        return self._simple_new(self.asi8, dtype=dtype, freq=self.freq)
 
     def tz_localize(self, tz, ambiguous='raise', nonexistent='raise',
                     errors=None):
@@ -995,7 +1011,8 @@ def tz_localize(self, tz, ambiguous='raise', nonexistent='raise',
                 self.asi8, tz, ambiguous=ambiguous, nonexistent=nonexistent,
             )
         new_dates = new_dates.view(_NS_DTYPE)
-        return self._simple_new(new_dates, tz=tz, freq=self.freq)
+        dtype = tz_to_dtype(tz)
+        return self._simple_new(new_dates, dtype=dtype, freq=self.freq)
 
     # ----------------------------------------------------------------
     # Conversion Methods - Vectorized analogues of Timestamp methods

pandas/core/indexes/datetimes.py

@@ -22,7 +22,7 @@
 
 from pandas.core.accessor import delegate_names
 from pandas.core.arrays.datetimes import (
-    DatetimeArray, _to_M8, validate_tz_from_dtype)
+    DatetimeArray, _to_M8, tz_to_dtype, validate_tz_from_dtype)
 from pandas.core.base import _shared_docs
 import pandas.core.common as com
 from pandas.core.indexes.base import Index
@@ -326,7 +326,9 @@ def _simple_new(cls, values, name=None, freq=None, tz=None, dtype=None):
         # DatetimeArray._simple_new will accept either i8 or M8[ns] dtypes
         if isinstance(values, DatetimeIndex):
             values = values._data
-        dtarr = DatetimeArray._simple_new(values, freq=freq, tz=tz)
+
+        dtype = tz_to_dtype(tz)
+        dtarr = DatetimeArray._simple_new(values, freq=freq, dtype=dtype)
         assert isinstance(dtarr, DatetimeArray)
 
         result = object.__new__(cls)
@@ -401,7 +403,8 @@ def __setstate__(self, state):
 
                 freq = own_state[1]
                 tz = timezones.tz_standardize(own_state[2])
-                dtarr = DatetimeArray._simple_new(data, freq=freq, tz=tz)
+                dtype = tz_to_dtype(tz)
+                dtarr = DatetimeArray._simple_new(data, freq=freq, dtype=dtype)
 
                 self.name = own_state[0]
 

pandas/core/internals/blocks.py

@@ -2389,7 +2389,7 @@ def _try_coerce_result(self, result):
                 result = result.reshape(np.prod(result.shape))
             # GH#24096 new values invalidates a frequency
             result = self._holder._simple_new(result, freq=None,
-                                              tz=self.values.tz)
+                                              dtype=self.values.dtype)
 
         return result
 

pandas/tseries/offsets.py

@@ -935,7 +935,7 @@ def apply_index(self, i):
         shifted = liboffsets.shift_months(i.asi8, self.n, self._day_opt)
         # TODO: going through __new__ raises on call to _validate_frequency;
         #  are we passing incorrect freq?
-        return type(i)._simple_new(shifted, freq=i.freq, tz=i.tz)
+        return type(i)._simple_new(shifted, freq=i.freq, dtype=i.dtype)
 
 
 class MonthEnd(MonthOffset):
@@ -1642,7 +1642,7 @@ def apply_index(self, dtindex):
         # TODO: going through __new__ raises on call to _validate_frequency;
         #  are we passing incorrect freq?
         return type(dtindex)._simple_new(shifted, freq=dtindex.freq,
-                                         tz=dtindex.tz)
+                                         dtype=dtindex.dtype)
 
 
 class BQuarterEnd(QuarterOffset):
@@ -1722,7 +1722,7 @@ def apply_index(self, dtindex):
         # TODO: going through __new__ raises on call to _validate_frequency;
         #  are we passing incorrect freq?
         return type(dtindex)._simple_new(shifted, freq=dtindex.freq,
-                                         tz=dtindex.tz)
+                                         dtype=dtindex.dtype)
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):