diff --git a/pandas/_libs/tslibs/dtypes.pyi b/pandas/_libs/tslibs/dtypes.pyi index edc3de05f5df2..5c343f89f38ea 100644 --- a/pandas/_libs/tslibs/dtypes.pyi +++ b/pandas/_libs/tslibs/dtypes.pyi @@ -56,3 +56,19 @@ class Resolution(Enum): def get_reso_from_freqstr(cls, freq: str) -> Resolution: ... @property def attr_abbrev(self) -> str: ... + +class NpyDatetimeUnit(Enum): + NPY_FR_Y: int + NPY_FR_M: int + NPY_FR_W: int + NPY_FR_D: int + NPY_FR_h: int + NPY_FR_m: int + NPY_FR_s: int + NPY_FR_ms: int + NPY_FR_us: int + NPY_FR_ns: int + NPY_FR_ps: int + NPY_FR_fs: int + NPY_FR_as: int + NPY_FR_GENERIC: int diff --git a/pandas/_libs/tslibs/dtypes.pyx b/pandas/_libs/tslibs/dtypes.pyx index 5da3944bfb147..e1640a074834e 100644 --- a/pandas/_libs/tslibs/dtypes.pyx +++ b/pandas/_libs/tslibs/dtypes.pyx @@ -255,6 +255,26 @@ class Resolution(Enum): return cls.from_attrname(attr_name) +class NpyDatetimeUnit(Enum): + """ + Python-space analogue to NPY_DATETIMEUNIT. + """ + NPY_FR_Y = NPY_DATETIMEUNIT.NPY_FR_Y + NPY_FR_M = NPY_DATETIMEUNIT.NPY_FR_M + NPY_FR_W = NPY_DATETIMEUNIT.NPY_FR_W + NPY_FR_D = NPY_DATETIMEUNIT.NPY_FR_D + NPY_FR_h = NPY_DATETIMEUNIT.NPY_FR_h + NPY_FR_m = NPY_DATETIMEUNIT.NPY_FR_m + NPY_FR_s = NPY_DATETIMEUNIT.NPY_FR_s + NPY_FR_ms = NPY_DATETIMEUNIT.NPY_FR_ms + NPY_FR_us = NPY_DATETIMEUNIT.NPY_FR_us + NPY_FR_ns = NPY_DATETIMEUNIT.NPY_FR_ns + NPY_FR_ps = NPY_DATETIMEUNIT.NPY_FR_ps + NPY_FR_fs = NPY_DATETIMEUNIT.NPY_FR_fs + NPY_FR_as = NPY_DATETIMEUNIT.NPY_FR_as + NPY_FR_GENERIC = NPY_DATETIMEUNIT.NPY_FR_GENERIC + + cdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit): if unit == NPY_DATETIMEUNIT.NPY_FR_ns or unit == NPY_DATETIMEUNIT.NPY_FR_GENERIC: # generic -> default to nanoseconds diff --git a/pandas/core/dtypes/dtypes.py b/pandas/core/dtypes/dtypes.py index 64d46976b54f6..90a22fc361ad3 100644 --- a/pandas/core/dtypes/dtypes.py +++ b/pandas/core/dtypes/dtypes.py @@ -670,14 +670,17 @@ class DatetimeTZDtype(PandasExtensionDtype): type: type[Timestamp] = Timestamp kind: str_type = "M" - str = "|M8[ns]" num = 101 - base = np.dtype("M8[ns]") + base = np.dtype("M8[ns]") # TODO: depend on reso? na_value = NaT _metadata = ("unit", "tz") _match = re.compile(r"(datetime64|M8)\[(?P.+), (?P.+)\]") _cache_dtypes: dict[str_type, PandasExtensionDtype] = {} + @cache_readonly + def str(self): + return f"|M8[{self._unit}]" + def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None: if isinstance(unit, DatetimeTZDtype): # error: "str" has no attribute "tz" @@ -696,8 +699,8 @@ def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None: "'DatetimeTZDtype.construct_from_string()' instead." ) raise ValueError(msg) - else: - raise ValueError("DatetimeTZDtype only supports ns units") + if unit not in ["s", "ms", "us", "ns"]: + raise ValueError("DatetimeTZDtype only supports s, ms, us, ns units") if tz: tz = timezones.maybe_get_tz(tz) @@ -710,6 +713,19 @@ def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None: self._unit = unit self._tz = tz + @cache_readonly + def _reso(self) -> int: + """ + The NPY_DATETIMEUNIT corresponding to this dtype's resolution. + """ + reso = { + "s": dtypes.NpyDatetimeUnit.NPY_FR_s, + "ms": dtypes.NpyDatetimeUnit.NPY_FR_ms, + "us": dtypes.NpyDatetimeUnit.NPY_FR_us, + "ns": dtypes.NpyDatetimeUnit.NPY_FR_ns, + }[self._unit] + return reso.value + @property def unit(self) -> str_type: """ diff --git a/pandas/tests/dtypes/test_dtypes.py b/pandas/tests/dtypes/test_dtypes.py index b7de8016f8fac..aef61045179ef 100644 --- a/pandas/tests/dtypes/test_dtypes.py +++ b/pandas/tests/dtypes/test_dtypes.py @@ -4,6 +4,8 @@ import pytest import pytz +from pandas._libs.tslibs.dtypes import NpyDatetimeUnit + from pandas.core.dtypes.base import _registry as registry from pandas.core.dtypes.common import ( is_bool_dtype, @@ -263,10 +265,17 @@ def test_hash_vs_equality(self, dtype): assert dtype2 != dtype4 assert hash(dtype2) != hash(dtype4) - def test_construction(self): - msg = "DatetimeTZDtype only supports ns units" + def test_construction_non_nanosecond(self): + res = DatetimeTZDtype("ms", "US/Eastern") + assert res.unit == "ms" + assert res._reso == NpyDatetimeUnit.NPY_FR_ms.value + assert res.str == "|M8[ms]" + assert str(res) == "datetime64[ms, US/Eastern]" + + def test_day_not_supported(self): + msg = "DatetimeTZDtype only supports s, ms, us, ns units" with pytest.raises(ValueError, match=msg): - DatetimeTZDtype("ms", "US/Eastern") + DatetimeTZDtype("D", "US/Eastern") def test_subclass(self): a = DatetimeTZDtype.construct_from_string("datetime64[ns, US/Eastern]")