Skip to content

ENH: TDA+datetime_scalar support non-nano #47675

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 12, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,6 @@
from pandas.util._exceptions import find_stack_level

from pandas.core.dtypes.common import (
DT64NS_DTYPE,
is_all_strings,
is_categorical_dtype,
is_datetime64_any_dtype,
Expand Down Expand Up @@ -1103,6 +1102,7 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
self = cast("TimedeltaArray", self)

from pandas.core.arrays import DatetimeArray
from pandas.core.arrays.datetimes import tz_to_dtype

assert other is not NaT
other = Timestamp(other)
Expand All @@ -1113,10 +1113,17 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
# Preserve our resolution
return DatetimeArray._simple_new(result, dtype=result.dtype)

if self._reso != other._reso:
raise NotImplementedError(
"Addition between TimedeltaArray and Timestamp with mis-matched "
"resolutions is not yet supported."
)

i8 = self.asi8
result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE
return DatetimeArray(result, dtype=dtype, freq=self.freq)
dtype = tz_to_dtype(tz=other.tz, unit=self._unit)
res_values = result.view(f"M8[{self._unit}]")
return DatetimeArray._simple_new(res_values, dtype=dtype, freq=self.freq)

@final
def _add_datetime_arraylike(self, other) -> DatetimeArray:
Expand Down
19 changes: 12 additions & 7 deletions pandas/core/arrays/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,22 +91,23 @@
_midnight = time(0, 0)


def tz_to_dtype(tz):
def tz_to_dtype(tz: tzinfo | None, unit: str = "ns"):
"""
Return a datetime64[ns] dtype appropriate for the given timezone.
Parameters
----------
tz : tzinfo or None
unit : str, default "ns"
Returns
-------
np.dtype or Datetime64TZDType
"""
if tz is None:
return DT64NS_DTYPE
return np.dtype(f"M8[{unit}]")
else:
return DatetimeTZDtype(tz=tz)
return DatetimeTZDtype(tz=tz, unit=unit)


def _field_accessor(name: str, field: str, docstring=None):
Expand Down Expand Up @@ -800,7 +801,7 @@ def tz_convert(self, tz) -> DatetimeArray:
)

# No conversion since timestamps are all UTC to begin with
dtype = tz_to_dtype(tz)
dtype = tz_to_dtype(tz, unit=self._unit)
return self._simple_new(self._ndarray, dtype=dtype, freq=self.freq)

@dtl.ravel_compat
Expand Down Expand Up @@ -965,10 +966,14 @@ def tz_localize(self, tz, ambiguous="raise", nonexistent="raise") -> DatetimeArr
# Convert to UTC

new_dates = tzconversion.tz_localize_to_utc(
self.asi8, tz, ambiguous=ambiguous, nonexistent=nonexistent
self.asi8,
tz,
ambiguous=ambiguous,
nonexistent=nonexistent,
reso=self._reso,
)
new_dates = new_dates.view(DT64NS_DTYPE)
dtype = tz_to_dtype(tz)
new_dates = new_dates.view(f"M8[{self._unit}]")
dtype = tz_to_dtype(tz, unit=self._unit)

freq = None
if timezones.is_utc(tz) or (len(self) == 1 and not isna(new_dates[0])):
Expand Down
28 changes: 28 additions & 0 deletions pandas/tests/arrays/test_timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,34 @@ def test_add_pdnat(self, tda):
assert result._reso == tda._reso
assert result.isna().all()

# TODO: 2022-07-11 this is the only test that gets to DTA.tz_convert
# or tz_localize with non-nano; implement tests specific to that.
def test_add_datetimelike_scalar(self, tda, tz_naive_fixture):
ts = pd.Timestamp("2016-01-01", tz=tz_naive_fixture)

msg = "with mis-matched resolutions"
with pytest.raises(NotImplementedError, match=msg):
# mismatched reso -> check that we don't give an incorrect result
tda + ts
with pytest.raises(NotImplementedError, match=msg):
# mismatched reso -> check that we don't give an incorrect result
ts + tda

ts = ts._as_unit(tda._unit)

exp_values = tda._ndarray + ts.asm8
expected = (
DatetimeArray._simple_new(exp_values, dtype=exp_values.dtype)
.tz_localize("UTC")
.tz_convert(ts.tz)
)

result = tda + ts
tm.assert_extension_array_equal(result, expected)

result = ts + tda
tm.assert_extension_array_equal(result, expected)

def test_mul_scalar(self, tda):
other = 2
result = tda * other
Expand Down