TYP: Add typing for pd.to_timedelta, fix Timedelta operator return types (#46343)

Author: Dr-Irv

pandas/_libs/tslibs/timedeltas.pyi

@@ -1,19 +1,65 @@
 from datetime import timedelta
 from typing import (
     ClassVar,
+    Literal,
     Type,
     TypeVar,
     overload,
 )
 
 import numpy as np
-import numpy.typing as npt
 
 from pandas._libs.tslibs import (
     NaTType,
     Tick,
 )
+from pandas._typing import npt
 
+# This should be kept consistent with the keys in the dict timedelta_abbrevs
+# in pandas/_libs/tslibs/timedeltas.pyx
+UnitChoices = Literal[
+    "Y",
+    "y",
+    "M",
+    "W",
+    "w",
+    "D",
+    "d",
+    "days",
+    "day",
+    "hours",
+    "hour",
+    "hr",
+    "h",
+    "m",
+    "minute",
+    "min",
+    "minutes",
+    "t",
+    "s",
+    "seconds",
+    "sec",
+    "second",
+    "ms",
+    "milliseconds",
+    "millisecond",
+    "milli",
+    "millis",
+    "l",
+    "us",
+    "microseconds",
+    "microsecond",
+    "µs",
+    "micro",
+    "micros",
+    "u",
+    "ns",
+    "nanoseconds",
+    "nano",
+    "nanos",
+    "nanosecond",
+    "n",
+]
 _S = TypeVar("_S", bound=timedelta)
 
 def ints_to_pytimedelta(
@@ -25,7 +71,7 @@ def array_to_timedelta64(
     unit: str | None = ...,
     errors: str = ...,
 ) -> np.ndarray: ...  # np.ndarray[m8ns]
-def parse_timedelta_unit(unit: str | None) -> str: ...
+def parse_timedelta_unit(unit: str | None) -> UnitChoices: ...
 def delta_to_nanoseconds(delta: np.timedelta64 | timedelta | Tick) -> int: ...
 
 class Timedelta(timedelta):
@@ -59,20 +105,20 @@ class Timedelta(timedelta):
     def ceil(self: _S, freq: str) -> _S: ...
     @property
     def resolution_string(self) -> str: ...
-    def __add__(self, other: timedelta) -> timedelta: ...
-    def __radd__(self, other: timedelta) -> timedelta: ...
-    def __sub__(self, other: timedelta) -> timedelta: ...
-    def __rsub__(self, other: timedelta) -> timedelta: ...
-    def __neg__(self) -> timedelta: ...
-    def __pos__(self) -> timedelta: ...
-    def __abs__(self) -> timedelta: ...
-    def __mul__(self, other: float) -> timedelta: ...
-    def __rmul__(self, other: float) -> timedelta: ...
+    def __add__(self, other: timedelta) -> Timedelta: ...
+    def __radd__(self, other: timedelta) -> Timedelta: ...
+    def __sub__(self, other: timedelta) -> Timedelta: ...
+    def __rsub__(self, other: timedelta) -> Timedelta: ...
+    def __neg__(self) -> Timedelta: ...
+    def __pos__(self) -> Timedelta: ...
+    def __abs__(self) -> Timedelta: ...
+    def __mul__(self, other: float) -> Timedelta: ...
+    def __rmul__(self, other: float) -> Timedelta: ...
     # error: Signature of "__floordiv__" incompatible with supertype "timedelta"
     @overload  # type: ignore[override]
     def __floordiv__(self, other: timedelta) -> int: ...
     @overload
-    def __floordiv__(self, other: int | float) -> timedelta: ...
+    def __floordiv__(self, other: int | float) -> Timedelta: ...
     @overload
     def __floordiv__(
         self, other: npt.NDArray[np.timedelta64]
@@ -90,11 +136,13 @@ class Timedelta(timedelta):
     @overload
     def __truediv__(self, other: timedelta) -> float: ...
     @overload
-    def __truediv__(self, other: float) -> timedelta: ...
-    def __mod__(self, other: timedelta) -> timedelta: ...
-    def __divmod__(self, other: timedelta) -> tuple[int, timedelta]: ...
+    def __truediv__(self, other: float) -> Timedelta: ...
+    def __mod__(self, other: timedelta) -> Timedelta: ...
+    def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ...
     def __le__(self, other: timedelta) -> bool: ...
     def __lt__(self, other: timedelta) -> bool: ...
     def __ge__(self, other: timedelta) -> bool: ...
     def __gt__(self, other: timedelta) -> bool: ...
     def __hash__(self) -> int: ...
+    def isoformat(self) -> str: ...
+    def to_numpy(self) -> np.timedelta64: ...

pandas/_libs/tslibs/timedeltas.pyx

@@ -82,6 +82,7 @@ Components = collections.namedtuple(
     ],
 )
 
+# This should be kept consistent with UnitChoices in pandas/_libs/tslibs/timedeltas.pyi
 cdef dict timedelta_abbrevs = {
     "Y": "Y",
     "y": "Y",

pandas/_typing.py

@@ -313,3 +313,4 @@ def closed(self) -> bool:
 # datetime and NaTType
 
 DatetimeNaTType = Union[datetime, "NaTType"]
+DateTimeErrorChoices = Literal["ignore", "raise", "coerce"]

pandas/core/tools/datetimes.py

@@ -39,6 +39,7 @@
 from pandas._typing import (
     AnyArrayLike,
     ArrayLike,
+    DateTimeErrorChoices,
     Timezone,
 )
 from pandas.util._exceptions import find_stack_level
@@ -79,6 +80,7 @@
 
 if TYPE_CHECKING:
     from pandas._libs.tslibs.nattype import NaTType
+    from pandas._libs.tslibs.timedeltas import UnitChoices
 
     from pandas import (
         DataFrame,
@@ -657,7 +659,7 @@ def _adjust_to_origin(arg, origin, unit):
 @overload
 def to_datetime(
     arg: DatetimeScalar,
-    errors: str = ...,
+    errors: DateTimeErrorChoices = ...,
     dayfirst: bool = ...,
     yearfirst: bool = ...,
     utc: bool | None = ...,
@@ -674,7 +676,7 @@ def to_datetime(
 @overload
 def to_datetime(
     arg: Series | DictConvertible,
-    errors: str = ...,
+    errors: DateTimeErrorChoices = ...,
     dayfirst: bool = ...,
     yearfirst: bool = ...,
     utc: bool | None = ...,
@@ -691,7 +693,7 @@ def to_datetime(
 @overload
 def to_datetime(
     arg: list | tuple | Index | ArrayLike,
-    errors: str = ...,
+    errors: DateTimeErrorChoices = ...,
     dayfirst: bool = ...,
     yearfirst: bool = ...,
     utc: bool | None = ...,
@@ -707,7 +709,7 @@ def to_datetime(
 
 def to_datetime(
     arg: DatetimeScalarOrArrayConvertible | DictConvertible,
-    errors: str = "raise",
+    errors: DateTimeErrorChoices = "raise",
     dayfirst: bool = False,
     yearfirst: bool = False,
     utc: bool | None = None,
@@ -1148,7 +1150,7 @@ def to_datetime(
 }
 
 
-def _assemble_from_unit_mappings(arg, errors, tz):
+def _assemble_from_unit_mappings(arg, errors: DateTimeErrorChoices, tz):
     """
     assemble the unit specified fields from the arg (DataFrame)
     Return a Series for actual parsing
@@ -1228,7 +1230,8 @@ def coerce(values):
     except (TypeError, ValueError) as err:
         raise ValueError(f"cannot assemble the datetimes: {err}") from err
 
-    for u in ["h", "m", "s", "ms", "us", "ns"]:
+    units: list[UnitChoices] = ["h", "m", "s", "ms", "us", "ns"]
+    for u in units:
         value = unit_rev.get(u)
         if value is not None and value in arg:
             try:

pandas/core/tools/timedeltas.py

@@ -3,6 +3,12 @@
 """
 from __future__ import annotations
 
+from datetime import timedelta
+from typing import (
+    TYPE_CHECKING,
+    overload,
+)
+
 import numpy as np
 
 from pandas._libs import lib
@@ -23,8 +29,61 @@
 
 from pandas.core.arrays.timedeltas import sequence_to_td64ns
 
-
-def to_timedelta(arg, unit=None, errors="raise"):
+if TYPE_CHECKING:
+    from pandas._libs.tslibs.timedeltas import UnitChoices
+    from pandas._typing import (
+        ArrayLike,
+        DateTimeErrorChoices,
+    )
+
+    from pandas import (
+        Index,
+        Series,
+        TimedeltaIndex,
+    )
+
+
+@overload
+def to_timedelta(
+    arg: str | int | float | timedelta,
+    unit: UnitChoices | None = ...,
+    errors: DateTimeErrorChoices = ...,
+) -> Timedelta:
+    ...
+
+
+@overload
+def to_timedelta(
+    arg: Series,
+    unit: UnitChoices | None = ...,
+    errors: DateTimeErrorChoices = ...,
+) -> Series:
+    ...
+
+
+@overload
+def to_timedelta(
+    arg: list | tuple | range | ArrayLike | Index,
+    unit: UnitChoices | None = ...,
+    errors: DateTimeErrorChoices = ...,
+) -> TimedeltaIndex:
+    ...
+
+
+def to_timedelta(
+    arg: str
+    | int
+    | float
+    | timedelta
+    | list
+    | tuple
+    | range
+    | ArrayLike
+    | Index
+    | Series,
+    unit: UnitChoices | None = None,
+    errors: DateTimeErrorChoices = "raise",
+) -> Timedelta | TimedeltaIndex | Series:
     """
     Convert argument to timedelta.
 
@@ -133,7 +192,11 @@ def to_timedelta(arg, unit=None, errors="raise"):
         return _convert_listlike(arg, unit=unit, errors=errors, name=arg.name)
     elif isinstance(arg, np.ndarray) and arg.ndim == 0:
         # extract array scalar and process below
-        arg = lib.item_from_zerodim(arg)
+        # error: Incompatible types in assignment (expression has type "object",
+        # variable has type "Union[str, int, float, timedelta, List[Any],
+        # Tuple[Any, ...], Union[Union[ExtensionArray, ndarray[Any, Any]], Index,
+        # Series]]")  [assignment]
+        arg = lib.item_from_zerodim(arg)  # type: ignore[assignment]
     elif is_list_like(arg) and getattr(arg, "ndim", 1) == 1:
         return _convert_listlike(arg, unit=unit, errors=errors)
     elif getattr(arg, "ndim", 1) > 1:

pandas/io/sql.py

@@ -25,7 +25,10 @@
 import numpy as np
 
 import pandas._libs.lib as lib
-from pandas._typing import DtypeArg
+from pandas._typing import (
+    DateTimeErrorChoices,
+    DtypeArg,
+)
 from pandas.compat._optional import import_optional_dependency
 from pandas.errors import AbstractMethodError
 from pandas.util._exceptions import find_stack_level
@@ -86,7 +89,7 @@ def _handle_date_column(
         # read_sql like functions.
         # Format can take on custom to_datetime argument values such as
         # {"errors": "coerce"} or {"dayfirst": True}
-        error = format.pop("errors", None) or "ignore"
+        error: DateTimeErrorChoices = format.pop("errors", None) or "ignore"
         return to_datetime(col, errors=error, **format)
     else:
         # Allow passing of formatting string for integers