Skip to content

Commit 259a15c

Browse files
authored
ENH: DatetimeTZDtype support non-nano (#47120)
* ENH: DatetimeTZDtype support non-nano * mypy fixup * mypy fixup * TST: Day not supported
1 parent a4c6a69 commit 259a15c

File tree

4 files changed

+68
-7
lines changed

4 files changed

+68
-7
lines changed

pandas/_libs/tslibs/dtypes.pyi

+16
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,19 @@ class Resolution(Enum):
5656
def get_reso_from_freqstr(cls, freq: str) -> Resolution: ...
5757
@property
5858
def attr_abbrev(self) -> str: ...
59+
60+
class NpyDatetimeUnit(Enum):
61+
NPY_FR_Y: int
62+
NPY_FR_M: int
63+
NPY_FR_W: int
64+
NPY_FR_D: int
65+
NPY_FR_h: int
66+
NPY_FR_m: int
67+
NPY_FR_s: int
68+
NPY_FR_ms: int
69+
NPY_FR_us: int
70+
NPY_FR_ns: int
71+
NPY_FR_ps: int
72+
NPY_FR_fs: int
73+
NPY_FR_as: int
74+
NPY_FR_GENERIC: int

pandas/_libs/tslibs/dtypes.pyx

+20
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,26 @@ class Resolution(Enum):
257257
return cls.from_attrname(attr_name)
258258

259259

260+
class NpyDatetimeUnit(Enum):
261+
"""
262+
Python-space analogue to NPY_DATETIMEUNIT.
263+
"""
264+
NPY_FR_Y = NPY_DATETIMEUNIT.NPY_FR_Y
265+
NPY_FR_M = NPY_DATETIMEUNIT.NPY_FR_M
266+
NPY_FR_W = NPY_DATETIMEUNIT.NPY_FR_W
267+
NPY_FR_D = NPY_DATETIMEUNIT.NPY_FR_D
268+
NPY_FR_h = NPY_DATETIMEUNIT.NPY_FR_h
269+
NPY_FR_m = NPY_DATETIMEUNIT.NPY_FR_m
270+
NPY_FR_s = NPY_DATETIMEUNIT.NPY_FR_s
271+
NPY_FR_ms = NPY_DATETIMEUNIT.NPY_FR_ms
272+
NPY_FR_us = NPY_DATETIMEUNIT.NPY_FR_us
273+
NPY_FR_ns = NPY_DATETIMEUNIT.NPY_FR_ns
274+
NPY_FR_ps = NPY_DATETIMEUNIT.NPY_FR_ps
275+
NPY_FR_fs = NPY_DATETIMEUNIT.NPY_FR_fs
276+
NPY_FR_as = NPY_DATETIMEUNIT.NPY_FR_as
277+
NPY_FR_GENERIC = NPY_DATETIMEUNIT.NPY_FR_GENERIC
278+
279+
260280
cdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit):
261281
if unit == NPY_DATETIMEUNIT.NPY_FR_ns or unit == NPY_DATETIMEUNIT.NPY_FR_GENERIC:
262282
# generic -> default to nanoseconds

pandas/core/dtypes/dtypes.py

+20-4
Original file line numberDiff line numberDiff line change
@@ -670,14 +670,17 @@ class DatetimeTZDtype(PandasExtensionDtype):
670670

671671
type: type[Timestamp] = Timestamp
672672
kind: str_type = "M"
673-
str = "|M8[ns]"
674673
num = 101
675-
base = np.dtype("M8[ns]")
674+
base = np.dtype("M8[ns]") # TODO: depend on reso?
676675
na_value = NaT
677676
_metadata = ("unit", "tz")
678677
_match = re.compile(r"(datetime64|M8)\[(?P<unit>.+), (?P<tz>.+)\]")
679678
_cache_dtypes: dict[str_type, PandasExtensionDtype] = {}
680679

680+
@cache_readonly
681+
def str(self):
682+
return f"|M8[{self._unit}]"
683+
681684
def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None:
682685
if isinstance(unit, DatetimeTZDtype):
683686
# error: "str" has no attribute "tz"
@@ -696,8 +699,8 @@ def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None:
696699
"'DatetimeTZDtype.construct_from_string()' instead."
697700
)
698701
raise ValueError(msg)
699-
else:
700-
raise ValueError("DatetimeTZDtype only supports ns units")
702+
if unit not in ["s", "ms", "us", "ns"]:
703+
raise ValueError("DatetimeTZDtype only supports s, ms, us, ns units")
701704

702705
if tz:
703706
tz = timezones.maybe_get_tz(tz)
@@ -710,6 +713,19 @@ def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None:
710713
self._unit = unit
711714
self._tz = tz
712715

716+
@cache_readonly
717+
def _reso(self) -> int:
718+
"""
719+
The NPY_DATETIMEUNIT corresponding to this dtype's resolution.
720+
"""
721+
reso = {
722+
"s": dtypes.NpyDatetimeUnit.NPY_FR_s,
723+
"ms": dtypes.NpyDatetimeUnit.NPY_FR_ms,
724+
"us": dtypes.NpyDatetimeUnit.NPY_FR_us,
725+
"ns": dtypes.NpyDatetimeUnit.NPY_FR_ns,
726+
}[self._unit]
727+
return reso.value
728+
713729
@property
714730
def unit(self) -> str_type:
715731
"""

pandas/tests/dtypes/test_dtypes.py

+12-3
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import pytest
55
import pytz
66

7+
from pandas._libs.tslibs.dtypes import NpyDatetimeUnit
8+
79
from pandas.core.dtypes.base import _registry as registry
810
from pandas.core.dtypes.common import (
911
is_bool_dtype,
@@ -263,10 +265,17 @@ def test_hash_vs_equality(self, dtype):
263265
assert dtype2 != dtype4
264266
assert hash(dtype2) != hash(dtype4)
265267

266-
def test_construction(self):
267-
msg = "DatetimeTZDtype only supports ns units"
268+
def test_construction_non_nanosecond(self):
269+
res = DatetimeTZDtype("ms", "US/Eastern")
270+
assert res.unit == "ms"
271+
assert res._reso == NpyDatetimeUnit.NPY_FR_ms.value
272+
assert res.str == "|M8[ms]"
273+
assert str(res) == "datetime64[ms, US/Eastern]"
274+
275+
def test_day_not_supported(self):
276+
msg = "DatetimeTZDtype only supports s, ms, us, ns units"
268277
with pytest.raises(ValueError, match=msg):
269-
DatetimeTZDtype("ms", "US/Eastern")
278+
DatetimeTZDtype("D", "US/Eastern")
270279

271280
def test_subclass(self):
272281
a = DatetimeTZDtype.construct_from_string("datetime64[ns, US/Eastern]")

0 commit comments

Comments
 (0)