TYP: unit str->TimeUnit (#62707)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -157,6 +157,8 @@
         Sequence,
     )
 
+    from pandas._typing import TimeUnit
+
     from pandas import Index
     from pandas.core.arrays import (
         DatetimeArray,
@@ -2114,7 +2116,7 @@ def _creso(self) -> int:
         return get_unit_from_dtype(self._ndarray.dtype)
 
     @cache_readonly
-    def unit(self) -> str:
+    def unit(self) -> TimeUnit:
         """
         The precision unit of the datetime data.
 
@@ -2138,11 +2140,11 @@ def unit(self) -> str:
         >>> idx.as_unit("s").unit
         's'
         """
-        # error: Argument 1 to "dtype_to_unit" has incompatible type
-        # "ExtensionDtype"; expected "Union[DatetimeTZDtype, dtype[Any]]"
-        return dtype_to_unit(self.dtype)  # type: ignore[arg-type]
+        # error: Incompatible return value type (got "str", expected
+        # "Literal['s', 'ms', 'us', 'ns']")  [return-value]
+        return dtype_to_unit(self.dtype)  # type: ignore[return-value,arg-type]
 
-    def as_unit(self, unit: str, round_ok: bool = True) -> Self:
+    def as_unit(self, unit: TimeUnit, round_ok: bool = True) -> Self:
         """
         Convert to a dtype with the given unit resolution.
 

pandas/core/arrays/datetimes.py

@@ -88,6 +88,7 @@
         IntervalClosedType,
         TimeAmbiguous,
         TimeNonexistent,
+        TimeUnit,
         npt,
     )
 
@@ -394,7 +395,10 @@ def _from_sequence_not_strict(
         result = cls._simple_new(subarr, freq=inferred_freq, dtype=data_dtype)
         if unit is not None and unit != result.unit:
             # If unit was specified in user-passed dtype, cast to it here
-            result = result.as_unit(unit)
+            # error: Argument 1 to "as_unit" of "TimelikeOps" has
+            # incompatible type "str"; expected "Literal['s', 'ms', 'us', 'ns']"
+            # [arg-type]
+            result = result.as_unit(unit)  # type: ignore[arg-type]
 
         validate_kwds = {"ambiguous": ambiguous}
         result._maybe_pin_freq(freq, validate_kwds)
@@ -413,7 +417,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]):
@@ -534,7 +538,7 @@ def _unbox_scalar(self, value) -> np.datetime64:
             raise ValueError("'value' should be a Timestamp.")
         self._check_compatible_with(value)
         if value is NaT:
-            return np.datetime64(value._value, self.unit)  # type: ignore[call-overload]
+            return np.datetime64(value._value, self.unit)
         else:
             return value.as_unit(self.unit, round_ok=False).asm8
 

pandas/core/arrays/period.py

@@ -962,7 +962,9 @@ def astype(self, dtype, copy: bool = True):
             # GH#45038 match PeriodIndex behavior.
             tz = getattr(dtype, "tz", None)
             unit = dtl.dtype_to_unit(dtype)
-            return self.to_timestamp().tz_localize(tz).as_unit(unit)
+            # error: Argument 1 to "as_unit" of "TimelikeOps" has incompatible
+            # type "str"; expected "Literal['s', 'ms', 'us', 'ns']"  [arg-type]
+            return self.to_timestamp().tz_localize(tz).as_unit(unit)  # type: ignore[arg-type]
 
         return super().astype(dtype, copy=copy)
 

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)
@@ -327,7 +325,7 @@ def _unbox_scalar(self, value) -> np.timedelta64:
             raise ValueError("'value' should be a Timedelta.")
         self._check_compatible_with(value)
         if value is NaT:
-            return np.timedelta64(value._value, self.unit)  # type: ignore[call-overload]
+            return np.timedelta64(value._value, self.unit)
         else:
             return value.as_unit(self.unit, round_ok=False).asm8
 

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
@@ -3094,7 +3095,7 @@ def asfreq(
 
         new_obj.index = _asfreq_compat(obj.index, freq)
     else:
-        unit = None
+        unit: TimeUnit = "ns"
         if isinstance(obj.index, DatetimeIndex):
             # TODO: should we disallow non-DatetimeIndex?
             unit = obj.index.unit