From bee17cd5fe3a254a9c282b7b51e69f80b0ded6ba Mon Sep 17 00:00:00 2001 From: Brock Mendel Date: Mon, 7 Jan 2019 15:11:51 -0800 Subject: [PATCH] Have DTA._simple_new take dtype instead of tz --- pandas/core/arrays/datetimelike.py | 17 +++++------- pandas/core/arrays/datetimes.py | 43 +++++++++++++++++++++--------- pandas/core/indexes/datetimes.py | 9 ++++--- pandas/core/internals/blocks.py | 2 +- pandas/tseries/offsets.py | 6 ++--- 5 files changed, 46 insertions(+), 31 deletions(-) diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 2de22e062b29b..cfb697b3c357a 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/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, diff --git a/pandas/core/arrays/datetimes.py b/pandas/core/arrays/datetimes.py index 708b4b074abcc..efa1757a989fc 100644 --- a/pandas/core/arrays/datetimes.py +++ b/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,8 +431,9 @@ 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: @@ -427,7 +441,8 @@ def _generate_range(cls, start, end, periods, freq, tz=None, 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 diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index 1f2a3636033b4..664ca9c5d2f05 100644 --- a/pandas/core/indexes/datetimes.py +++ b/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] diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 13dca433f9ead..70e4f44cb5de8 100644 --- a/pandas/core/internals/blocks.py +++ b/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 diff --git a/pandas/tseries/offsets.py b/pandas/tseries/offsets.py index 573f02fe0aa52..73f85d954432e 100644 --- a/pandas/tseries/offsets.py +++ b/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):