TYP: stronger typing for unit, as_unit

Author: jbrockmendel

pandas/_libs/tslibs/nattype.pyi

@@ -17,6 +17,7 @@ from pandas._libs.tslibs.period import Period
 from pandas._typing import (
     Frequency,
     TimestampNonexistent,
+    TimeUnit,
 )
 
 NaT: NaTType
@@ -180,4 +181,4 @@ class NaTType:
     def __floordiv__(self, other: float, /) -> Self: ...
     # other
     def __hash__(self) -> int: ...
-    def as_unit(self, unit: str, round_ok: bool = ...) -> NaTType: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> NaTType: ...

pandas/_libs/tslibs/timedeltas.pyi

@@ -15,6 +15,7 @@ from pandas._libs.tslibs import (
 )
 from pandas._typing import (
     Frequency,
+    TimeUnit,
     npt,
 )
 
@@ -162,5 +163,5 @@ class Timedelta(timedelta):
     ) -> np.timedelta64: ...
     def view(self, dtype: npt.DTypeLike) -> object: ...
     @property
-    def unit(self) -> str: ...
-    def as_unit(self, unit: str, round_ok: bool = ...) -> Timedelta: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Timedelta: ...

pandas/_libs/tslibs/timestamps.pyi

@@ -23,7 +23,10 @@ from pandas._libs.tslibs import (
     Tick,
     Timedelta,
 )
-from pandas._typing import TimestampNonexistent
+from pandas._typing import (
+    TimestampNonexistent,
+    TimeUnit,
+)
 
 _TimeZones: TypeAlias = str | _tzinfo | None | int
 
@@ -235,5 +238,5 @@ class Timestamp(datetime):
     @property
     def daysinmonth(self) -> int: ...
     @property
-    def unit(self) -> str: ...
-    def as_unit(self, unit: str, round_ok: bool = ...) -> Timestamp: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Timestamp: ...

pandas/core/dtypes/dtypes.py

@@ -85,6 +85,7 @@
         IntervalClosedType,
         Ordered,
         Scalar,
+        TimeUnit,
         npt,
         type_t,
     )
@@ -820,7 +821,7 @@ def _creso(self) -> int:
         return abbrev_to_npy_unit(self.unit)
 
     @property
-    def unit(self) -> str_type:
+    def unit(self) -> TimeUnit:
         """
         The precision of the datetime data.
 

pandas/core/window/rolling.py

@@ -15,6 +15,7 @@
     Concatenate,
     Literal,
     Self,
+    cast,
     final,
     overload,
 )
@@ -95,6 +96,7 @@
         NDFrameT,
         QuantileInterpolation,
         P,
+        TimeUnit,
         T,
         WindowingRankType,
         npt,
@@ -2001,6 +2003,7 @@ def _validate(self) -> None:
                 except TypeError:
                     # if not a datetime dtype, eg for empty dataframes
                     unit = "ns"
+                unit = cast("TimeUnit", unit)
                 self._win_freq_i8 = Timedelta(freq.nanos).as_unit(unit)._value
 
             # min_periods must be an integer