ENH: Improve general datetime functions (#352)

* TYP: Improve infer_freq

* TYP: Imrprove timedelta_range

* TYP: Imrprove period_range

* TYP: Improve date_range and bdate_range

* TYP: Improve to_timedelta

* TYP: Improve to_datetime

* Small fixes

* TST: Add tests for to_timedelta

* TST/ENH: Add tests and overload for bdate_range

* TST/ENH: Add tests and improve period_range

* TST: Add tests

* TST: Add check and assert type, add more tests

* CLN: Remove unnecessary TODO

* ENH: Improve timestamp convertible

Co-authored-by: Kevin Sheppard <[email protected]>

Author: bashtage

pandas-stubs/_libs/tslibs/timedeltas.pyi

@@ -3,6 +3,7 @@ from typing import (
     ClassVar,
     Literal,
     TypeVar,
+    Union,
     overload,
 )
 
@@ -17,10 +18,7 @@ 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: TypeAlias = Literal[
-    "Y",
-    "y",
-    "M",
+TimeDeltaUnitChoices: TypeAlias = Literal[
     "W",
     "w",
     "D",
@@ -60,6 +58,16 @@ UnitChoices: TypeAlias = Literal[
     "nanosecond",
     "n",
 ]
+
+UnitChoices: TypeAlias = Union[
+    TimeDeltaUnitChoices,
+    Literal[
+        "Y",
+        "y",
+        "M",
+    ],
+]
+
 _S = TypeVar("_S", bound=timedelta)
 
 def ints_to_pytimedelta(

pandas-stubs/_typing.pyi

@@ -19,6 +19,7 @@ from typing import (
 import numpy as np
 from numpy import typing as npt
 from pandas.core.arrays import ExtensionArray
+from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
 from pandas.core.groupby.grouper import Grouper
 from pandas.core.indexes.base import Index
@@ -46,6 +47,27 @@ PandasScalar: TypeAlias = Union[
 
 DatetimeLike: TypeAlias = Union[datetime.datetime, np.datetime64, Timestamp]
 
+DatetimeDictArg: TypeAlias = Union[
+    Sequence[int], Sequence[float], list[str], tuple[Scalar, ...], AnyArrayLike
+]
+DictConvertible: TypeAlias = Union[FulldatetimeDict, DataFrame]
+
+class YearMonthDayDict(TypedDict, total=True):
+    year: DatetimeDictArg
+    month: DatetimeDictArg
+    day: DatetimeDictArg
+
+class FulldatetimeDict(YearMonthDayDict, total=False):
+    hour: DatetimeDictArg
+    hours: DatetimeDictArg
+    minute: DatetimeDictArg
+    minutes: DatetimeDictArg
+    second: DatetimeDictArg
+    seconds: DatetimeDictArg
+    ms: DatetimeDictArg
+    us: DatetimeDictArg
+    ns: DatetimeDictArg
+
 # dtypes
 NpDtype: TypeAlias = Union[
     str, np.dtype[np.generic], type[Union[str, complex, bool, object]]
@@ -178,7 +200,7 @@ IndexingInt: TypeAlias = Union[
     int, np.int_, np.integer, np.unsignedinteger, np.signedinteger, np.int8
 ]
 TimestampConvertibleTypes: TypeAlias = Union[
-    Timestamp, datetime.datetime, np.datetime64, np.int64, float, str
+    Timestamp, datetime.datetime, datetime.date, np.datetime64, np.int64, float, str
 ]
 TimedeltaConvertibleTypes: TypeAlias = Union[
     Timedelta, datetime.timedelta, np.timedelta64, np.int64, float, str

pandas-stubs/core/indexes/datetimes.pyi

@@ -1,8 +1,13 @@
 from datetime import (
+    date,
     timedelta,
     tzinfo,
 )
-from typing import overload
+from typing import (
+    Hashable,
+    Sequence,
+    overload,
+)
 
 import numpy as np
 from pandas import (
@@ -90,20 +95,33 @@ def date_range(
     freq: str | timedelta | Timedelta | BaseOffset = ...,
     tz: str | tzinfo = ...,
     normalize: bool = ...,
-    name: str | None = ...,
+    name: Hashable | None = ...,
     inclusive: IntervalClosedType = ...,
-    **kwargs,
 ) -> DatetimeIndex: ...
+@overload
 def bdate_range(
     start: str | DatetimeLike | None = ...,
     end: str | DatetimeLike | None = ...,
     periods: int | None = ...,
     freq: str | timedelta | Timedelta | BaseOffset = ...,
     tz: str | tzinfo = ...,
     normalize: bool = ...,
-    name: str | None = ...,
+    name: Hashable | None = ...,
+    weekmask: str | None = ...,
+    holidays: None = ...,
+    inclusive: IntervalClosedType = ...,
+) -> DatetimeIndex: ...
+@overload
+def bdate_range(
+    start: str | DatetimeLike | None = ...,
+    end: str | DatetimeLike | None = ...,
+    periods: int | None = ...,
+    *,
+    freq: str | timedelta | Timedelta | BaseOffset,
+    tz: str | tzinfo = ...,
+    normalize: bool = ...,
+    name: Hashable | None = ...,
     weekmask: str | None = ...,
-    holidays: list | None = ...,
+    holidays: Sequence[str | DatetimeLike | date],
     inclusive: IntervalClosedType = ...,
-    **kwargs,
 ) -> DatetimeIndex: ...

pandas-stubs/core/indexes/period.pyi

@@ -1,9 +1,14 @@
+from typing import Hashable
+
 import numpy as np
+import pandas as pd
 from pandas.core.indexes.datetimelike import (
     DatetimeIndexOpsMixin as DatetimeIndexOpsMixin,
 )
 from pandas.core.indexes.numeric import Int64Index
 
+from pandas._libs.tslibs import BaseOffset
+
 class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
     def __new__(
         cls,
@@ -46,5 +51,9 @@ class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
     def memory_usage(self, deep: bool = ...): ...
 
 def period_range(
-    start=..., end=..., periods=..., freq=..., name=...
+    start: str | pd.Period | None = ...,
+    end: str | pd.Period | None = ...,
+    periods: int | None = ...,
+    freq: str | BaseOffset | None = ...,
+    name: Hashable | None = ...,
 ) -> PeriodIndex: ...

pandas-stubs/core/indexes/timedeltas.pyi

@@ -1,5 +1,10 @@
-from typing import overload
+from typing import (
+    Hashable,
+    Literal,
+    overload,
+)
 
+from pandas import DateOffset
 from pandas.core.indexes.accessors import TimedeltaIndexProperties
 from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
 from pandas.core.indexes.datetimes import DatetimeIndex
@@ -9,7 +14,10 @@ from pandas._libs import (
     Timedelta,
     Timestamp,
 )
-from pandas._typing import num
+from pandas._typing import (
+    TimedeltaConvertibleTypes,
+    num,
+)
 
 class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties):
     def __new__(
@@ -42,5 +50,10 @@ class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties):
     def to_series(self, index=..., name=...) -> TimedeltaSeries: ...
 
 def timedelta_range(
-    start=..., end=..., periods=..., freq=..., name=..., closed=...
+    start: TimedeltaConvertibleTypes = ...,
+    end: TimedeltaConvertibleTypes = ...,
+    periods: int | None = ...,
+    freq: str | DateOffset | None = ...,
+    name: Hashable | None = ...,
+    closed: Literal["left", "right"] | None = ...,
 ) -> TimedeltaIndex: ...

pandas-stubs/core/tools/datetimes.pyi

@@ -1,18 +1,21 @@
-from datetime import datetime
+from datetime import (
+    date,
+    datetime,
+)
 from typing import (
     Literal,
-    TypedDict,
+    Sequence,
     Union,
     overload,
 )
 
 import numpy as np
+import pandas as pd
 from pandas import (
     Index,
     Timestamp,
 )
 from pandas.core.arrays import ExtensionArray
-from pandas.core.frame import DataFrame
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.core.series import (
     Series,
@@ -24,38 +27,18 @@ from pandas._libs.tslibs import NaTType
 from pandas._typing import (
     AnyArrayLike,
     DateTimeErrorChoices,
+    DictConvertible,
     IgnoreRaise,
+    TimestampConvertibleTypes,
+    npt,
 )
 
 ArrayConvertible: TypeAlias = Union[list, tuple, AnyArrayLike]
 Scalar: TypeAlias = Union[float, str]
-DatetimeScalar: TypeAlias = Union[Scalar, datetime, np.datetime64]
+DatetimeScalar: TypeAlias = Union[Scalar, datetime, np.datetime64, date]
 
 DatetimeScalarOrArrayConvertible: TypeAlias = Union[DatetimeScalar, ArrayConvertible]
 
-DatetimeDictArg: TypeAlias = Union[list[Scalar], tuple[Scalar, ...], AnyArrayLike]
-
-class YearMonthDayDict(TypedDict, total=True):
-    year: DatetimeDictArg
-    month: DatetimeDictArg
-    day: DatetimeDictArg
-
-class FulldatetimeDict(YearMonthDayDict, total=False):
-    hour: DatetimeDictArg
-    hours: DatetimeDictArg
-    minute: DatetimeDictArg
-    minutes: DatetimeDictArg
-    second: DatetimeDictArg
-    seconds: DatetimeDictArg
-    ms: DatetimeDictArg
-    us: DatetimeDictArg
-    ns: DatetimeDictArg
-
-DictConvertible: TypeAlias = Union[FulldatetimeDict, DataFrame]
-
-def should_cache(
-    arg: ArrayConvertible, unique_share: float = ..., check_count: int | None = ...
-) -> bool: ...
 @overload
 def to_datetime(
     arg: DatetimeScalar,
@@ -67,7 +50,7 @@ def to_datetime(
     exact: bool = ...,
     unit: str | None = ...,
     infer_datetime_format: bool = ...,
-    origin=...,
+    origin: Literal["julian", "unix"] | TimestampConvertibleTypes = ...,
     cache: bool = ...,
 ) -> Timestamp: ...
 @overload
@@ -81,7 +64,7 @@ def to_datetime(
     exact: bool = ...,
     unit: str | None = ...,
     infer_datetime_format: bool = ...,
-    origin=...,
+    origin: Literal["julian", "unix"] | TimestampConvertibleTypes = ...,
     cache: bool = ...,
 ) -> Timestamp | NaTType: ...
 @overload
@@ -95,12 +78,19 @@ def to_datetime(
     exact: bool = ...,
     unit: str | None = ...,
     infer_datetime_format: bool = ...,
-    origin=...,
+    origin: Literal["julian", "unix"] | TimestampConvertibleTypes = ...,
     cache: bool = ...,
 ) -> TimestampSeries: ...
 @overload
 def to_datetime(
-    arg: list | tuple | np.ndarray | Index | ExtensionArray,
+    arg: Sequence[int | float | datetime]
+    | list[str]
+    | tuple[int | float | str | datetime, ...]
+    | npt.NDArray[np.datetime64]
+    | npt.NDArray[np.str_]
+    | npt.NDArray[np.int_]
+    | Index
+    | ExtensionArray,
     errors: DateTimeErrorChoices = ...,
     dayfirst: bool = ...,
     yearfirst: bool = ...,
@@ -109,7 +99,6 @@ def to_datetime(
     exact: bool = ...,
     unit: str | None = ...,
     infer_datetime_format: bool = ...,
-    origin=...,
+    origin: Literal["julian", "unix"] | TimestampConvertibleTypes = ...,
     cache: bool = ...,
 ) -> DatetimeIndex: ...
-def to_time(arg, format=..., infer_time_format: bool = ..., errors: str = ...): ...

pandas-stubs/core/tools/timedeltas.pyi

@@ -1,7 +1,10 @@
-# def to_timedelta(arg, unit: str = ..., errors: str = ...): ...
 from datetime import timedelta
-from typing import overload
+from typing import (
+    Sequence,
+    overload,
+)
 
+import pandas as pd
 from pandas import Index
 from pandas.core.indexes.timedeltas import TimedeltaIndex
 from pandas.core.series import (
@@ -10,29 +13,32 @@ from pandas.core.series import (
 )
 
 from pandas._libs.tslibs import Timedelta
-from pandas._libs.tslibs.timedeltas import UnitChoices
+from pandas._libs.tslibs.timedeltas import TimeDeltaUnitChoices
 from pandas._typing import (
     ArrayLike,
     DateTimeErrorChoices,
 )
 
-# Copied from pandas/_libs/tslibs/timedeltas.pyx
-
 @overload
 def to_timedelta(
     arg: str | float | timedelta,
-    unit: UnitChoices | None = ...,
+    unit: TimeDeltaUnitChoices | None = ...,
     errors: DateTimeErrorChoices = ...,
 ) -> Timedelta: ...
 @overload
 def to_timedelta(
     arg: Series,
-    unit: UnitChoices | None = ...,
+    unit: TimeDeltaUnitChoices | None = ...,
     errors: DateTimeErrorChoices = ...,
 ) -> TimedeltaSeries: ...
 @overload
 def to_timedelta(
-    arg: list | tuple | range | ArrayLike | Index,
-    unit: UnitChoices | None = ...,
+    arg: Sequence[float | timedelta]
+    | list[str | float | timedelta]
+    | tuple[str | float | timedelta, ...]
+    | range
+    | ArrayLike
+    | Index,
+    unit: TimeDeltaUnitChoices | None = ...,
     errors: DateTimeErrorChoices = ...,
 ) -> TimedeltaIndex: ...

pandas-stubs/tseries/frequencies.pyi

@@ -1,6 +1,12 @@
+from pandas import (
+    DatetimeIndex,
+    Series,
+    TimedeltaIndex,
+)
+
 from pandas.tseries.offsets import DateOffset as DateOffset
 
 def get_period_alias(offset_str: str) -> str | None: ...
 def to_offset(freq) -> DateOffset | None: ...
 def get_offset(name: str) -> DateOffset: ...
-def infer_freq(index) -> str | None: ...
+def infer_freq(index: Series | DatetimeIndex | TimedeltaIndex) -> str | None: ...