TYP: unit str->TimeUnit

Author: jbrockmendel

pandas/core/arrays/datetimes.py

@@ -88,6 +88,7 @@
         IntervalClosedType,
         TimeAmbiguous,
         TimeNonexistent,
+        TimeUnit,
         npt,
     )
 
@@ -413,7 +414,7 @@ def _generate_range(
         nonexistent: TimeNonexistent = "raise",
         inclusive: IntervalClosedType = "both",
         *,
-        unit: str | None = None,
+        unit: TimeUnit = "ns",
     ) -> Self:
         periods = dtl.validate_periods(periods)
         if freq is None and any(x is None for x in [periods, start, end]):

pandas/core/arrays/timedeltas.py

@@ -73,6 +73,7 @@
         DtypeObj,
         NpDtype,
         npt,
+        TimeUnit,
     )
 
     from pandas import DataFrame
@@ -275,7 +276,7 @@ def _from_sequence_not_strict(
 
     @classmethod
     def _generate_range(
-        cls, start, end, periods, freq, closed=None, *, unit: str | None = None
+        cls, start, end, periods, freq, closed=None, *, unit: TimeUnit
     ) -> Self:
         periods = dtl.validate_periods(periods)
         if freq is None and any(x is None for x in [periods, start, end]):
@@ -293,11 +294,8 @@ def _generate_range(
         if end is not None:
             end = Timedelta(end).as_unit("ns")
 
-        if unit is not None:
-            if unit not in ["s", "ms", "us", "ns"]:
-                raise ValueError("'unit' must be one of 's', 'ms', 'us', 'ns'")
-        else:
-            unit = "ns"
+        if unit not in ["s", "ms", "us", "ns"]:
+            raise ValueError("'unit' must be one of 's', 'ms', 'us', 'ns'")
 
         if start is not None and unit is not None:
             start = start.as_unit(unit, round_ok=False)

pandas/core/indexes/datetimelike.py

@@ -75,6 +75,7 @@
     from pandas._typing import (
         Axis,
         JoinHow,
+        TimeUnit,
         npt,
     )
 
@@ -434,10 +435,10 @@ class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin, ABC):
     _is_unique = Index.is_unique
 
     @property
-    def unit(self) -> str:
+    def unit(self) -> TimeUnit:
         return self._data.unit
 
-    def as_unit(self, unit: str) -> Self:
+    def as_unit(self, unit: TimeUnit) -> Self:
         """
         Convert to a dtype with the given unit resolution.
 

pandas/core/indexes/datetimes.py

@@ -66,6 +66,7 @@
         TimeAmbiguous,
         TimeNonexistent,
         npt,
+        TimeUnit,
     )
 
     from pandas.core.api import (
@@ -852,7 +853,7 @@ def date_range(
     name: Hashable | None = None,
     inclusive: IntervalClosedType = "both",
     *,
-    unit: str | None = None,
+    unit: TimeUnit = "ns",
     **kwargs,
 ) -> DatetimeIndex:
     """
@@ -893,7 +894,7 @@ def date_range(
         Include boundaries; Whether to set each bound as closed or open.
 
         .. versionadded:: 1.4.0
-    unit : str, default None
+    unit : {'s', 'ms', 'us', 'ns'}, default 'ns'
         Specify the desired resolution of the result.
 
         .. versionadded:: 2.0.0

pandas/core/indexes/timedeltas.py

@@ -33,7 +33,10 @@
 
 if TYPE_CHECKING:
     from pandas._libs import NaTType
-    from pandas._typing import DtypeObj
+    from pandas._typing import (
+        DtypeObj,
+        TimeUnit,
+    )
 
 
 @inherit_names(
@@ -249,7 +252,7 @@ def timedelta_range(
     name=None,
     closed=None,
     *,
-    unit: str | None = None,
+    unit: TimeUnit = "ns",
 ) -> TimedeltaIndex:
     """
     Return a fixed frequency TimedeltaIndex with day as the default.
@@ -269,7 +272,7 @@ def timedelta_range(
     closed : str, default None
         Make the interval closed with respect to the given frequency to
         the 'left', 'right', or both sides (None).
-    unit : str, default None
+    unit : {'s', 'ms', 'us', 'ns'}, default 'ns'
         Specify the desired resolution of the result.
 
         .. versionadded:: 2.0.0

pandas/core/resample.py

@@ -99,6 +99,7 @@
         TimedeltaConvertibleTypes,
         TimeGrouperOrigin,
         TimestampConvertibleTypes,
+        TimeUnit,
         npt,
     )
 
@@ -2835,7 +2836,7 @@ def _get_timestamp_range_edges(
     first: Timestamp,
     last: Timestamp,
     freq: BaseOffset,
-    unit: str,
+    unit: TimeUnit,
     closed: Literal["right", "left"] = "left",
     origin: TimeGrouperOrigin = "start_day",
     offset: Timedelta | None = None,
@@ -2985,7 +2986,7 @@ def _adjust_dates_anchored(
     closed: Literal["right", "left"] = "right",
     origin: TimeGrouperOrigin = "start_day",
     offset: Timedelta | None = None,
-    unit: str = "ns",
+    unit: TimeUnit = "ns",
 ) -> tuple[Timestamp, Timestamp]:
     # First and last offsets should be calculated from the start day to fix an
     # error cause by resampling across multiple days when a one day period is