cast to TimeUnit

Author: jbrockmendel

pandas/core/arrays/_ranges.py

@@ -22,15 +22,18 @@
 from pandas.core.construction import range_to_ndarray
 
 if TYPE_CHECKING:
-    from pandas._typing import npt
+    from pandas._typing import (
+        TimeUnit,
+        npt,
+    )
 
 
 def generate_regular_range(
     start: Timestamp | Timedelta | None,
     end: Timestamp | Timedelta | None,
     periods: int | None,
     freq: BaseOffset,
-    unit: str = "ns",
+    unit: TimeUnit = "ns",
 ) -> npt.NDArray[np.intp]:
     """
     Generate a range of dates or timestamps with the spans between dates
@@ -46,7 +49,7 @@ def generate_regular_range(
         Number of periods in produced date range.
     freq : Tick
         Describes space between dates in produced date range.
-    unit : str, default "ns"
+    unit : {'s', 'ms', 'us', 'ns'}, default "ns"
         The resolution the output is meant to represent.
 
     Returns

pandas/core/arrays/datetimes.py

@@ -106,15 +106,15 @@
 
 
 @overload
-def tz_to_dtype(tz: tzinfo, unit: str = ...) -> DatetimeTZDtype: ...
+def tz_to_dtype(tz: tzinfo, unit: TimeUnit = ...) -> DatetimeTZDtype: ...
 
 
 @overload
-def tz_to_dtype(tz: None, unit: str = ...) -> np.dtype[np.datetime64]: ...
+def tz_to_dtype(tz: None, unit: TimeUnit = ...) -> np.dtype[np.datetime64]: ...
 
 
 def tz_to_dtype(
-    tz: tzinfo | None, unit: str = "ns"
+    tz: tzinfo | None, unit: TimeUnit = "ns"
 ) -> np.dtype[np.datetime64] | DatetimeTZDtype:
     """
     Return a datetime64[ns] dtype appropriate for the given timezone.
@@ -393,6 +393,7 @@ def _from_sequence_not_strict(
             )
 
         data_unit = np.datetime_data(subarr.dtype)[0]
+        data_unit = cast("TimeUnit", data_unit)
         data_dtype = tz_to_dtype(tz, data_unit)
         result = cls._simple_new(subarr, freq=inferred_freq, dtype=data_dtype)
         if unit is not None and unit != result.unit:
@@ -2935,7 +2936,7 @@ def _generate_range(
     periods: int | None,
     offset: BaseOffset,
     *,
-    unit: str,
+    unit: TimeUnit,
 ) -> Generator[Timestamp]:
     """
     Generates a sequence of dates corresponding to the specified time

pandas/core/dtypes/cast.py

@@ -100,6 +100,7 @@
         DtypeObj,
         NumpyIndexT,
         Scalar,
+        TimeUnit,
     )
 
     from pandas import Index
@@ -567,6 +568,7 @@ def _maybe_promote(dtype: np.dtype, fill_value=np.nan):
             # different unit, e.g. passed np.timedelta64(24, "h") with dtype=m8[ns]
             # see if we can losslessly cast it to our dtype
             unit = np.datetime_data(dtype)[0]
+            unit = cast("TimeUnit", unit)
             try:
                 td = Timedelta(fill_value).as_unit(unit, round_ok=False)
             except OutOfBoundsTimedelta:

pandas/core/dtypes/dtypes.py

@@ -781,10 +781,10 @@ def base(self) -> DtypeObj:  # type: ignore[override]
     def str(self) -> str:  # type: ignore[override]
         return f"|M8[{self.unit}]"
 
-    def __init__(self, unit: str_type | DatetimeTZDtype = "ns", tz=None) -> None:
+    def __init__(self, unit: TimeUnit | DatetimeTZDtype = "ns", tz=None) -> None:
         if isinstance(unit, DatetimeTZDtype):
             # error: "str" has no attribute "tz"
-            unit, tz = unit.unit, unit.tz  # type: ignore[attr-defined]
+            unit, tz = unit.unit, unit.tz  # type: ignore[union-attr]
 
         if unit != "ns":
             if isinstance(unit, str) and tz is None:
@@ -895,7 +895,8 @@ def construct_from_string(cls, string: str_type) -> DatetimeTZDtype:
         if match:
             d = match.groupdict()
             try:
-                return cls(unit=d["unit"], tz=d["tz"])
+                unit = cast("TimeUnit", d["unit"])
+                return cls(unit=unit, tz=d["tz"])
             except (KeyError, TypeError, ValueError) as err:
                 # KeyError if maybe_get_tz tries and fails to get a
                 #  zoneinfo timezone (actually zoneinfo.ZoneInfoNotFoundError).
@@ -972,6 +973,7 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
         if all(isinstance(t, DatetimeTZDtype) and t.tz == self.tz for t in dtypes):
             np_dtype = np.max([cast(DatetimeTZDtype, t).base for t in [self, *dtypes]])
             unit = np.datetime_data(np_dtype)[0]
+            unit = cast("TimeUnit", unit)
             return type(self)(unit=unit, tz=self.tz)
         return super()._get_common_dtype(dtypes)
 

pandas/core/reshape/tile.py

@@ -8,6 +8,7 @@
     TYPE_CHECKING,
     Any,
     Literal,
+    cast,
 )
 
 import numpy as np
@@ -49,6 +50,7 @@
     from pandas._typing import (
         DtypeObj,
         IntervalLeftRight,
+        TimeUnit,
     )
 
 
@@ -412,7 +414,7 @@ def _nbins_to_bins(x_idx: Index, nbins: int, right: bool) -> Index:
             # error: Argument 1 to "dtype_to_unit" has incompatible type
             # "dtype[Any] | ExtensionDtype"; expected "DatetimeTZDtype | dtype[Any]"
             unit = dtype_to_unit(x_idx.dtype)  # type: ignore[arg-type]
-            td = Timedelta(seconds=1).as_unit(unit)
+            td = Timedelta(seconds=1).as_unit(cast("TimeUnit", unit))
             # Use DatetimeArray/TimedeltaArray method instead of linspace
             # error: Item "ExtensionArray" of "ExtensionArray | ndarray[Any, Any]"
             # has no attribute "_generate_range"
@@ -595,6 +597,7 @@ def _format_labels(
         # error: Argument 1 to "dtype_to_unit" has incompatible type
         # "dtype[Any] | ExtensionDtype"; expected "DatetimeTZDtype | dtype[Any]"
         unit = dtype_to_unit(bins.dtype)  # type: ignore[arg-type]
+        unit = cast("TimeUnit", unit)
         formatter = lambda x: x
         adjust = lambda x: x - Timedelta(1, unit=unit).as_unit(unit)
     else:

pandas/core/tools/datetimes.py

@@ -87,6 +87,7 @@
 
     from pandas._libs.tslibs.nattype import NaTType
     from pandas._libs.tslibs.timedeltas import UnitChoices
+    from pandas._typing import TimeUnit
 
     from pandas import (
         DataFrame,
@@ -447,6 +448,7 @@ def _convert_listlike_datetimes(
         # We can take a shortcut since the datetime64 numpy array
         # is in UTC
         out_unit = np.datetime_data(result.dtype)[0]
+        out_unit = cast("TimeUnit", out_unit)
         dtype = tz_to_dtype(tz_parsed, out_unit)
         dt64_values = result.view(f"M8[{dtype.unit}]")
         dta = DatetimeArray._simple_new(dt64_values, dtype=dtype)
@@ -469,13 +471,15 @@ def _array_strptime_with_fallback(
     result, tz_out = array_strptime(arg, fmt, exact=exact, errors=errors, utc=utc)
     if tz_out is not None:
         unit = np.datetime_data(result.dtype)[0]
+        unit = cast("TimeUnit", unit)
         dtype = DatetimeTZDtype(tz=tz_out, unit=unit)
         dta = DatetimeArray._simple_new(result, dtype=dtype)
         if utc:
             dta = dta.tz_convert("UTC")
         return Index(dta, name=name)
     elif result.dtype != object and utc:
         unit = np.datetime_data(result.dtype)[0]
+        unit = cast("TimeUnit", unit)
         res = Index(result, dtype=f"M8[{unit}, UTC]", name=name)
         return res
     return Index(result, dtype=result.dtype, name=name)

pandas/core/window/ewm.py

@@ -3,7 +3,10 @@
 import datetime
 from functools import partial
 from textwrap import dedent
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    cast,
+)
 
 import numpy as np
 
@@ -57,6 +60,7 @@
 if TYPE_CHECKING:
     from pandas._typing import (
         TimedeltaConvertibleTypes,
+        TimeUnit,
         npt,
     )
 
@@ -122,6 +126,7 @@ def _calculate_deltas(
         Diff of the times divided by the half-life
     """
     unit = dtype_to_unit(times.dtype)
+    unit = cast("TimeUnit", unit)
     if isinstance(times, ABCSeries):
         times = times._values
     _times = np.asarray(times.view(np.int64), dtype=np.float64)

pandas/io/pytables.py

@@ -134,6 +134,7 @@
         AxisInt,
         DtypeArg,
         FilePath,
+        TimeUnit,
         npt,
     )
 
@@ -5093,6 +5094,9 @@ def _set_tz(
     # Argument "tz" to "tz_to_dtype" has incompatible type "str | tzinfo | None";
     # expected "tzinfo"
     unit, _ = np.datetime_data(datetime64_dtype)  # parsing dtype: unit, count
+    unit = cast("TimeUnit", unit)
+    # error: Argument "tz" to "tz_to_dtype" has incompatible type
+    #  "str | tzinfo | None"; expected "tzinfo"
     dtype = tz_to_dtype(tz=tz, unit=unit)  # type: ignore[arg-type]
     dta = DatetimeArray._from_sequence(values, dtype=dtype)
     return dta