diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index fdb0a398c..f22a12778 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -21,7 +21,7 @@ from pandas.core.series import ( from pandas._typing import ( IntervalClosedType, IntervalT, - np_ndarray_bool, + np_1darray, npt, ) @@ -170,7 +170,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __gt__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __gt__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __gt__( self, @@ -179,7 +181,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __lt__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __lt__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __lt__( self, @@ -188,7 +192,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __ge__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __ge__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __ge__( self, @@ -197,11 +203,15 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __le__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __le__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __eq__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __eq__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload @@ -209,7 +219,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __ne__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... + def __ne__( + self: IntervalT, other: IntervalIndex[IntervalT] + ) -> np_1darray[np.bool]: ... @overload def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index d13c4675c..603ae742b 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -22,7 +22,11 @@ from typing_extensions import TypeAlias from pandas._libs.tslibs import NaTType from pandas._libs.tslibs.offsets import BaseOffset from pandas._libs.tslibs.timestamps import Timestamp -from pandas._typing import npt +from pandas._typing import ( + ShapeT, + np_1darray, + np_ndarray, +) class IncompatibleFrequency(ValueError): ... @@ -98,44 +102,64 @@ class Period(PeriodMixin): @overload def __eq__(self, other: Period) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... # type: ignore[overload-overlap] + def __eq__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] @overload def __eq__(self, other: PeriodSeries) -> Series[bool]: ... # type: ignore[overload-overlap] @overload + def __eq__(self, other: np_ndarray[ShapeT, np.object_]) -> np_ndarray[ShapeT, np.bool]: ... # type: ignore[overload-overlap] + @overload def __eq__(self, other: object) -> Literal[False]: ... @overload def __ge__(self, other: Period) -> bool: ... @overload - def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + def __ge__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... @overload def __ge__(self, other: PeriodSeries) -> Series[bool]: ... @overload + def __ge__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.bool]: ... + @overload def __gt__(self, other: Period) -> bool: ... @overload - def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + def __gt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... @overload def __gt__(self, other: PeriodSeries) -> Series[bool]: ... @overload + def __gt__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.bool]: ... + @overload def __le__(self, other: Period) -> bool: ... @overload - def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + def __le__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... @overload def __le__(self, other: PeriodSeries) -> Series[bool]: ... @overload + def __le__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.bool]: ... + @overload def __lt__(self, other: Period) -> bool: ... @overload - def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + def __lt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... @overload def __lt__(self, other: PeriodSeries) -> Series[bool]: ... + @overload + def __lt__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.bool]: ... # ignore[misc] here because we know all other comparisons # are False, so we use Literal[False] @overload def __ne__(self, other: Period) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... # type: ignore[overload-overlap] + def __ne__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] @overload def __ne__(self, other: PeriodSeries) -> Series[bool]: ... # type: ignore[overload-overlap] @overload + def __ne__(self, other: np_ndarray[ShapeT, np.object_]) -> np_ndarray[ShapeT, np.bool]: ... # type: ignore[overload-overlap] + @overload def __ne__(self, other: object) -> Literal[True]: ... # Ignored due to indecipherable error from mypy: # Forward operator "__add__" is not callable [misc] diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 8d5ce244b..54c0509ef 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -33,7 +33,10 @@ from pandas._libs.tslibs import ( from pandas._libs.tslibs.period import Period from pandas._libs.tslibs.timestamps import Timestamp from pandas._typing import ( + ShapeT, TimeUnit, + np_1darray, + np_ndarray, npt, ) @@ -153,12 +156,12 @@ class Timedelta(timedelta): def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ... @overload def __add__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __add__( - self, other: npt.NDArray[np.datetime64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... @overload def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... @overload @@ -176,12 +179,12 @@ class Timedelta(timedelta): def __radd__(self, other: NaTType) -> NaTType: ... @overload def __radd__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __radd__( - self, other: npt.NDArray[np.datetime64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... @overload def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... @overload @@ -193,8 +196,8 @@ class Timedelta(timedelta): def __sub__(self, other: NaTType) -> NaTType: ... @overload def __sub__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ... @overload @@ -215,12 +218,12 @@ class Timedelta(timedelta): def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ... @overload def __rsub__( - self, other: npt.NDArray[np.datetime64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... @overload def __rsub__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... def __neg__(self) -> Timedelta: ... @@ -231,8 +234,8 @@ class Timedelta(timedelta): def __mul__(self, other: float) -> Timedelta: ... @overload def __mul__( - self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __mul__(self, other: Series[int]) -> TimedeltaSeries: ... @overload @@ -243,8 +246,8 @@ class Timedelta(timedelta): def __rmul__(self, other: float) -> Timedelta: ... @overload def __rmul__( - self, other: npt.NDArray[np.floating] | npt.NDArray[np.integer] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.floating] | np_ndarray[ShapeT, np.integer] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ... @overload @@ -260,12 +263,12 @@ class Timedelta(timedelta): def __floordiv__(self, other: float) -> Timedelta: ... @overload def __floordiv__( - self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __floordiv__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.int_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.int_]: ... @overload def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... @overload @@ -282,8 +285,8 @@ class Timedelta(timedelta): def __rfloordiv__(self, other: NaTType | None) -> float: ... @overload def __rfloordiv__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.int_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.int_]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ... @@ -291,8 +294,8 @@ class Timedelta(timedelta): def __truediv__(self, other: float) -> Timedelta: ... @overload def __truediv__( - self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ... @overload @@ -308,9 +311,11 @@ class Timedelta(timedelta): @overload def __eq__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload + def __eq__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + @overload def __eq__( # type: ignore[overload-overlap] - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __eq__(self, other: object) -> Literal[False]: ... # Override due to more types supported than dt.timedelta @@ -319,9 +324,11 @@ class Timedelta(timedelta): @overload def __ne__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload + def __ne__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + @overload def __ne__( # type: ignore[overload-overlap] - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __ne__(self, other: object) -> Literal[True]: ... # Override due to more types supported than dt.timedelta @@ -335,8 +342,8 @@ class Timedelta(timedelta): def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... @overload def __mod__( - self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] - ) -> npt.NDArray[np.timedelta64]: ... + self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating] + ) -> np_ndarray[ShapeT, np.timedelta64]: ... @overload def __mod__( self, other: Series[int] | Series[float] | TimedeltaSeries @@ -348,36 +355,44 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload + def __le__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + @overload def __le__( - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __le__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload + def __lt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + @overload def __lt__( - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __lt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload + def __ge__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + @overload def __ge__( - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __ge__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload + def __gt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + @overload def __gt__( - self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.bool_]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.bool_]: ... @overload def __gt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... def __hash__(self) -> int: ... diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index 11302a12c..6146d0b36 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -40,10 +40,11 @@ from pandas._libs.tslibs import ( Timedelta, ) from pandas._typing import ( + ShapeT, TimestampNonexistent, TimeUnit, - np_ndarray_bool, - npt, + np_1darray, + np_ndarray, ) _Ambiguous: TypeAlias = bool | Literal["raise", "NaT"] @@ -180,40 +181,48 @@ class Timestamp(datetime, SupportsIndex): @overload # type: ignore[override] def __le__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload + def __le__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + @overload def __le__( - self, other: DatetimeIndex | npt.NDArray[np.datetime64] - ) -> np_ndarray_bool: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.bool]: ... @overload def __le__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] def __lt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload + def __lt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + @overload def __lt__( - self, other: DatetimeIndex | npt.NDArray[np.datetime64] - ) -> np_ndarray_bool: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.bool]: ... @overload def __lt__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] def __ge__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload + def __ge__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + @overload def __ge__( - self, other: DatetimeIndex | npt.NDArray[np.datetime64] - ) -> np_ndarray_bool: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.bool]: ... @overload def __ge__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] def __gt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload + def __gt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + @overload def __gt__( - self, other: DatetimeIndex | npt.NDArray[np.datetime64] - ) -> np_ndarray_bool: ... + self, other: np_ndarray[ShapeT, np.datetime64] + ) -> np_ndarray[ShapeT, np.bool]: ... @overload def __gt__(self, other: TimestampSeries) -> Series[bool]: ... # error: Signature of "__add__" incompatible with supertype "date"/"datetime" @overload # type: ignore[override] def __add__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... @overload def __add__(self, other: timedelta | np.timedelta64 | Tick) -> Self: ... @overload @@ -226,8 +235,8 @@ class Timestamp(datetime, SupportsIndex): def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ... @overload def __radd__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... # TODO: test dt64 @overload # type: ignore[override] def __sub__(self, other: Timestamp | datetime | np.datetime64) -> Timedelta: ... @@ -241,14 +250,16 @@ class Timestamp(datetime, SupportsIndex): def __sub__(self, other: TimestampSeries) -> TimedeltaSeries: ... @overload def __sub__( - self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.datetime64]: ... + self, other: np_ndarray[ShapeT, np.timedelta64] + ) -> np_ndarray[ShapeT, np.datetime64]: ... @overload def __eq__(self, other: Timestamp | datetime | np.datetime64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __eq__(self, other: TimestampSeries) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... # type: ignore[overload-overlap] + def __eq__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + @overload + def __eq__(self, other: np_ndarray[ShapeT, np.datetime64]) -> np_ndarray[ShapeT, np.bool]: ... # type: ignore[overload-overlap] @overload def __eq__(self, other: object) -> Literal[False]: ... @overload @@ -256,7 +267,9 @@ class Timestamp(datetime, SupportsIndex): @overload def __ne__(self, other: TimestampSeries) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... # type: ignore[overload-overlap] + def __ne__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + @overload + def __ne__(self, other: np_ndarray[ShapeT, np.datetime64]) -> np_ndarray[ShapeT, np.bool]: ... # type: ignore[overload-overlap] @overload def __ne__(self, other: object) -> Literal[True]: ... def __hash__(self) -> int: ... diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index f14c393fe..1dc22c5d9 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -819,6 +819,23 @@ np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating] np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_] np_ndarray_str: TypeAlias = npt.NDArray[np.str_] +# Define shape and generic type variables with defaults similar to numpy +GenericT = TypeVar("GenericT", bound=np.generic, default=Any) +GenericT_co = TypeVar("GenericT_co", bound=np.generic, default=Any, covariant=True) +ShapeT = TypeVar("ShapeT", bound=tuple[int, ...], default=tuple[Any, ...]) +# Numpy ndarray with more ergonomic typevar +np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]] +# Numpy arrays with known shape (Do not use as argument types, only as return types) +np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]] +np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] + +class SupportsDType(Protocol[GenericT_co]): + @property + def dtype(self) -> np.dtype[GenericT_co]: ... + +# Similar to npt.DTypeLike but leaves out np.dtype and None for use in overloads +DTypeLike: TypeAlias = type[Any] | tuple[Any, Any] | list[Any] | str + IndexType: TypeAlias = slice | np_ndarray_anyint | Index | list[int] | Series[int] MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool] diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 803878ba8..a89c4d829 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -18,6 +18,7 @@ from pandas._typing import ( AnyArrayLike, IntervalT, TakeIndexer, + np_1darray, ) # These are type: ignored because the Index types overlap due to inheritance but indices @@ -54,14 +55,14 @@ def factorize( sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np.ndarray, Index]: ... +) -> tuple[np_1darray, Index]: ... @overload def factorize( values: Categorical, sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np.ndarray, Categorical]: ... +) -> tuple[np_1darray, Categorical]: ... def value_counts( values: AnyArrayLike | list | tuple, sort: bool = True, diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index b5c0ef7cc..b59ef8639 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -12,6 +12,7 @@ from pandas._typing import ( ScalarIndexer, SequenceIndexer, TakeIndexer, + np_1darray, npt, ) @@ -31,7 +32,7 @@ class ExtensionArray: dtype: npt.DTypeLike | None = ..., copy: bool = False, na_value: Scalar = ..., - ) -> np.ndarray: ... + ) -> np_1darray: ... @property def dtype(self) -> ExtensionDtype: ... @property @@ -44,13 +45,13 @@ class ExtensionArray: def isna(self) -> ArrayLike: ... def argsort( self, *, ascending: bool = ..., kind: str = ..., **kwargs - ) -> np.ndarray: ... + ) -> np_1darray: ... def fillna(self, value=..., method=None, limit=None): ... def dropna(self): ... def shift(self, periods: int = 1, fill_value: object = ...) -> Self: ... def unique(self): ... def searchsorted(self, value, side: str = ..., sorter=...): ... - def factorize(self, use_na_sentinel: bool = True) -> tuple[np.ndarray, Self]: ... + def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ... def repeat(self, repeats, axis=...): ... def take( self, @@ -60,7 +61,7 @@ class ExtensionArray: fill_value=..., ) -> Self: ... def copy(self) -> Self: ... - def view(self, dtype=...) -> Self | np.ndarray: ... + def view(self, dtype=...) -> Self | np_1darray: ... def ravel(self, order="C") -> Self: ... def tolist(self) -> list: ... def _reduce( diff --git a/pandas-stubs/core/arrays/categorical.pyi b/pandas-stubs/core/arrays/categorical.pyi index 57f075e21..a0ef49f7b 100644 --- a/pandas-stubs/core/arrays/categorical.pyi +++ b/pandas-stubs/core/arrays/categorical.pyi @@ -25,8 +25,7 @@ from pandas._typing import ( ScalarIndexer, SequenceIndexer, TakeIndexer, - np_ndarray_bool, - np_ndarray_int, + np_1darray, ) from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype @@ -63,7 +62,7 @@ class Categorical(ExtensionArray): fastpath: bool = ..., ) -> Categorical: ... @property - def codes(self) -> np_ndarray_int: ... + def codes(self) -> np_1darray[np.signedinteger]: ... def set_ordered(self, value) -> Categorical: ... def as_ordered(self) -> Categorical: ... def as_unordered(self) -> Categorical: ... @@ -90,7 +89,7 @@ class Categorical(ExtensionArray): @property def shape(self): ... def shift(self, periods=1, fill_value=...): ... - def __array__(self, dtype=...) -> np.ndarray: ... + def __array__(self, dtype=...) -> np_1darray: ... def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ... @property def T(self): ... @@ -98,10 +97,10 @@ class Categorical(ExtensionArray): def nbytes(self) -> int: ... def memory_usage(self, deep: bool = ...): ... def searchsorted(self, value, side: str = ..., sorter=...): ... - def isna(self) -> np_ndarray_bool: ... - def isnull(self) -> np_ndarray_bool: ... - def notna(self) -> np_ndarray_bool: ... - def notnull(self) -> np_ndarray_bool: ... + def isna(self) -> np_1darray[np.bool]: ... + def isnull(self) -> np_1darray[np.bool]: ... + def notna(self) -> np_1darray[np.bool]: ... + def notnull(self) -> np_1darray[np.bool]: ... def dropna(self): ... def value_counts(self, dropna: bool = True): ... def check_for_ordered(self, op) -> None: ... diff --git a/pandas-stubs/core/arrays/interval.pyi b/pandas-stubs/core/arrays/interval.pyi index 304cc2960..a7c8d5a6e 100644 --- a/pandas-stubs/core/arrays/interval.pyi +++ b/pandas-stubs/core/arrays/interval.pyi @@ -21,7 +21,7 @@ from pandas._typing import ( ScalarIndexer, SequenceIndexer, TakeIndexer, - np_ndarray_bool, + np_1darray, ) IntervalOrNA: TypeAlias = Interval | float @@ -99,7 +99,7 @@ class IntervalArray(IntervalMixin, ExtensionArray): def mid(self) -> Index: ... @property def is_non_overlapping_monotonic(self) -> bool: ... - def __array__(self, dtype=...) -> np.ndarray: ... + def __array__(self, dtype=...) -> np_1darray: ... def __arrow_array__(self, type=...): ... def to_tuples(self, na_tuple: bool = True): ... def repeat(self, repeats, axis: Axis | None = ...): ... @@ -108,5 +108,5 @@ class IntervalArray(IntervalMixin, ExtensionArray): @overload def contains( self, other: Scalar | ExtensionArray | Index | np.ndarray - ) -> np_ndarray_bool: ... + ) -> np_1darray[np.bool]: ... def overlaps(self, other: Interval) -> bool: ... diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index da7b69b74..c7241a3a7 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -24,9 +24,13 @@ from pandas._typing import ( S1, AxisIndex, DropKeep, + DTypeLike, + GenericT, + GenericT_co, NDFrameT, Scalar, - npt, + SupportsDType, + np_1darray, ) from pandas.util._decorators import cache_readonly @@ -42,7 +46,7 @@ class SelectionMixin(Generic[NDFrameT]): def __getitem__(self, key): ... def aggregate(self, func, *args, **kwargs): ... -class IndexOpsMixin(OpsMixin, Generic[S1]): +class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]): __array_priority__: int = ... @property def T(self) -> Self: ... @@ -57,13 +61,30 @@ class IndexOpsMixin(OpsMixin, Generic[S1]): def size(self) -> int: ... @property def array(self) -> ExtensionArray: ... + @overload + def to_numpy( + self, + dtype: None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT_co]: ... + @overload + def to_numpy( + self, + dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT], + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT]: ... + @overload def to_numpy( self, - dtype: npt.DTypeLike | None = ..., + dtype: DTypeLike, copy: bool = False, na_value: Scalar = ..., **kwargs, - ) -> np.ndarray: ... + ) -> np_1darray: ... @property def empty(self) -> bool: ... def max(self, axis=..., skipna: bool = ..., **kwargs): ... @@ -114,7 +135,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1]): def is_monotonic_increasing(self) -> bool: ... def factorize( self, sort: bool = False, use_na_sentinel: bool = True - ) -> tuple[np.ndarray, np.ndarray | Index | Categorical]: ... + ) -> tuple[np_1darray, np_1darray | Index | Categorical]: ... def searchsorted( self, value, side: Literal["left", "right"] = ..., sorter=... ) -> int | list[int]: ... diff --git a/pandas-stubs/core/dtypes/missing.pyi b/pandas-stubs/core/dtypes/missing.pyi index e36496cfd..d53c0cf53 100644 --- a/pandas-stubs/core/dtypes/missing.pyi +++ b/pandas-stubs/core/dtypes/missing.pyi @@ -4,20 +4,23 @@ from typing import ( ) import numpy as np -from numpy import typing as npt from pandas import ( DataFrame, Index, Series, ) +from pandas.core.arrays import ExtensionArray from typing_extensions import TypeIs from pandas._libs.missing import NAType from pandas._libs.tslibs import NaTType from pandas._typing import ( - ArrayLike, Scalar, ScalarT, + ShapeT, + np_1darray, + np_ndarray, + np_ndarray_bool, ) isposinf_scalar = ... @@ -28,7 +31,11 @@ def isna(obj: DataFrame) -> DataFrame: ... @overload def isna(obj: Series) -> Series[bool]: ... @overload -def isna(obj: Index | list[Any] | ArrayLike) -> npt.NDArray[np.bool_]: ... +def isna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray[np.bool]: ... +@overload +def isna(obj: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ... +@overload +def isna(obj: list[Any]) -> np_ndarray_bool: ... @overload def isna( obj: Scalar | NaTType | NAType | None, @@ -41,7 +48,11 @@ def notna(obj: DataFrame) -> DataFrame: ... @overload def notna(obj: Series) -> Series[bool]: ... @overload -def notna(obj: Index | list[Any] | ArrayLike) -> npt.NDArray[np.bool_]: ... +def notna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray[np.bool]: ... +@overload +def notna(obj: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ... +@overload +def notna(obj: list[Any]) -> np_ndarray_bool: ... @overload def notna(obj: ScalarT | NaTType | NAType | None) -> TypeIs[ScalarT]: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index a33321d75..14d8bfea1 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -165,6 +165,7 @@ from pandas._typing import ( ValueKeyFunc, WriteBuffer, XMLParsers, + np_2darray, npt, num, ) @@ -452,7 +453,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): dtype: npt.DTypeLike | None = ..., copy: bool = False, na_value: Scalar = ..., - ) -> np.ndarray: ... + ) -> np_2darray: ... @overload def to_dict( self, @@ -1766,7 +1767,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @property def size(self) -> int: ... @property - def values(self) -> np.ndarray: ... + def values(self) -> np_2darray: ... # methods @final def abs(self) -> Self: ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index 6746c9dd4..e96615dee 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -10,7 +10,6 @@ from typing import ( ) import numpy as np -import numpy.typing as npt from pandas import ( DatetimeIndex, Index, @@ -40,6 +39,7 @@ from pandas._typing import ( TimestampConvention, TimeUnit, TimeZones, + np_1darray, np_ndarray_bool, ) @@ -88,7 +88,7 @@ class _DatetimeFieldOps( ): ... _DTBoolOpsReturnType = TypeVar( - "_DTBoolOpsReturnType", bound=Series[bool] | np_ndarray_bool + "_DTBoolOpsReturnType", bound=Series[bool] | np_1darray[np.bool] ) class _IsLeapYearProperty(Generic[_DTBoolOpsReturnType]): @@ -126,10 +126,10 @@ class _DatetimeObjectOps( ): ... _DTOtherOpsDateReturnType = TypeVar( - "_DTOtherOpsDateReturnType", bound=Series[dt.date] | np.ndarray + "_DTOtherOpsDateReturnType", bound=Series[dt.date] | np_1darray[np.object_] ) _DTOtherOpsTimeReturnType = TypeVar( - "_DTOtherOpsTimeReturnType", bound=Series[dt.time] | np.ndarray + "_DTOtherOpsTimeReturnType", bound=Series[dt.time] | np_1darray[np.object_] ) class _DatetimeOtherOps(Generic[_DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType]): @@ -280,7 +280,7 @@ class DatetimeProperties( _DTToPeriodReturnType, ], ): - def to_pydatetime(self) -> np.ndarray: ... + def to_pydatetime(self) -> np_1darray[np.object_]: ... def isocalendar(self) -> DataFrame: ... @property def unit(self) -> TimeUnit: ... @@ -296,7 +296,7 @@ _TDTotalSecondsReturnType = TypeVar( class _TimedeltaPropertiesNoRounding( Generic[_TDNoRoundingMethodReturnType, _TDTotalSecondsReturnType] ): - def to_pytimedelta(self) -> np.ndarray: ... + def to_pytimedelta(self) -> np_1darray[np.object_]: ... @property def components(self) -> DataFrame: ... @property @@ -401,10 +401,10 @@ class DatetimeIndexProperties( Properties, _DatetimeNoTZProperties[ Index[int], - np_ndarray_bool, + np_1darray[np.bool], DatetimeIndex, - np.ndarray, - np.ndarray, + np_1darray[np.object_], + np_1darray[np.object_], BaseOffset, DatetimeIndex, Index, @@ -416,7 +416,7 @@ class DatetimeIndexProperties( def is_normalized(self) -> bool: ... @property def tzinfo(self) -> _tzinfo | None: ... - def to_pydatetime(self) -> npt.NDArray[np.object_]: ... + def to_pydatetime(self) -> np_1darray[np.object_]: ... def std( self, axis: int | None = ..., ddof: int = ..., skipna: bool = ... ) -> Timedelta: ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 0d7b440e6..313c3bde7 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -13,10 +13,12 @@ from datetime import ( from typing import ( Any, ClassVar, + Generic, Literal, TypeAlias, final, overload, + type_check_only, ) import numpy as np @@ -49,7 +51,10 @@ from pandas._typing import ( AxesData, DropKeep, DtypeArg, + DTypeLike, DtypeObj, + GenericT, + GenericT_co, HashableT, IgnoreRaise, Label, @@ -60,10 +65,11 @@ from pandas._typing import ( Scalar, SequenceNotStr, SliceType, + SupportsDType, TimedeltaDtypeArg, TimestampDtypeArg, + np_1darray, np_ndarray_anyint, - np_ndarray_bool, np_ndarray_complex, np_ndarray_float, type_t, @@ -274,7 +280,7 @@ class Index(IndexOpsMixin[S1]): ) -> StringMethods[ Self, MultiIndex, - np_ndarray_bool, + np_1darray[np.bool], Index[list[_str]], Index[int], Index[bytes], @@ -286,7 +292,7 @@ class Index(IndexOpsMixin[S1]): def __len__(self) -> int: ... def __array__( self, dtype: _str | np.dtype = ..., copy: bool | None = ... - ) -> np.ndarray: ... + ) -> np_1darray: ... def __array_wrap__(self, result, context=...): ... @property def dtype(self) -> DtypeObj: ... @@ -354,7 +360,7 @@ class Index(IndexOpsMixin[S1]): def dropna(self, how: AnyAll = "any") -> Self: ... def unique(self, level=...) -> Self: ... def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ... - def duplicated(self, keep: DropKeep = "first") -> np_ndarray_bool: ... + def duplicated(self, keep: DropKeep = "first") -> np_1darray[np.bool]: ... def __and__(self, other: Never) -> Never: ... def __rand__(self, other: Never) -> Never: ... def __or__(self, other: Never) -> Never: ... @@ -378,7 +384,7 @@ class Index(IndexOpsMixin[S1]): result_name: Hashable = ..., sort: bool | None = None, ) -> Self: ... - def get_loc(self, key: Label) -> int | slice | np_ndarray_bool: ... + def get_loc(self, key: Label) -> int | slice | np_1darray[np.bool]: ... def get_indexer( self, target, method: ReindexMethod | None = ..., limit=..., tolerance=... ): ... @@ -400,7 +406,7 @@ class Index(IndexOpsMixin[S1]): sort: bool = ..., ): ... @property - def values(self) -> np.ndarray: ... + def values(self) -> np_1darray: ... @property def array(self) -> ExtensionArray: ... def memory_usage(self, deep: bool = False): ... @@ -445,7 +451,7 @@ class Index(IndexOpsMixin[S1]): @final def groupby(self, values) -> dict[Hashable, np.ndarray]: ... def map(self, mapper, na_action=...) -> Index: ... - def isin(self, values, level=...) -> np_ndarray_bool: ... + def isin(self, values, level=...) -> np_1darray[np.bool]: ... def slice_indexer( self, start: Label | None = None, @@ -462,13 +468,13 @@ class Index(IndexOpsMixin[S1]): @property def shape(self) -> tuple[int, ...]: ... # Extra methods from old stubs - def __eq__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __eq__(self, other: object) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] def __iter__(self) -> Iterator[S1]: ... - def __ne__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __le__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __ge__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __lt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __gt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __ne__(self, other: object) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __le__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __ge__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __lt__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __gt__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] # overwrite inherited methods from OpsMixin @overload def __mul__( @@ -519,3 +525,30 @@ class Index(IndexOpsMixin[S1]): def infer_objects(self, copy: bool = True) -> Self: ... UnknownIndex: TypeAlias = Index[Any] + +@type_check_only +class _IndexSubclassBase(Index[S1], Generic[S1, GenericT_co]): + @overload + def to_numpy( # pyrefly: ignore + self, + dtype: None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT_co]: ... + @overload + def to_numpy( + self, + dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT], + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT]: ... + @overload + def to_numpy( + self, + dtype: DTypeLike, + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray: ... diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index 1e14614c7..f3f86dedd 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -7,10 +7,11 @@ from pandas._libs.tslibs import BaseOffset from pandas._typing import ( S1, AxisIndex, + GenericT_co, TimeUnit, ) -class DatetimeIndexOpsMixin(ExtensionIndex[S1]): +class DatetimeIndexOpsMixin(ExtensionIndex[S1, GenericT_co]): @property def freq(self) -> BaseOffset | None: ... @property @@ -33,7 +34,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]): self, other: DatetimeIndexOpsMixin ) -> TimedeltaIndex: ... -class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): +class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1, GenericT_co]): @property def unit(self) -> TimeUnit: ... def as_unit(self, unit: TimeUnit) -> Self: ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index c79856ff7..8bcf703f7 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -43,7 +43,9 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype from pandas.tseries.offsets import BaseOffset -class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties): +class DatetimeIndex( + DatetimeTimedeltaMixin[Timestamp, np.datetime64], DatetimeIndexProperties +): def __new__( cls, data: AxesData, diff --git a/pandas-stubs/core/indexes/extension.pyi b/pandas-stubs/core/indexes/extension.pyi index 293784afa..418816d8d 100644 --- a/pandas-stubs/core/indexes/extension.pyi +++ b/pandas-stubs/core/indexes/extension.pyi @@ -1,5 +1,8 @@ -from pandas.core.indexes.base import Index +from pandas.core.indexes.base import _IndexSubclassBase -from pandas._typing import S1 +from pandas._typing import ( + S1, + GenericT_co, +) -class ExtensionIndex(Index[S1]): ... +class ExtensionIndex(_IndexSubclassBase[S1, GenericT_co]): ... diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index b82cdeddb..b59a47bc4 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -32,6 +32,7 @@ from pandas._typing import ( IntervalT, Label, MaskType, + np_1darray, np_ndarray_anyint, np_ndarray_bool, npt, @@ -65,7 +66,7 @@ _EdgesTimedelta: TypeAlias = ( _TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime _TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta -class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): +class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): closed: IntervalClosedType def __new__( @@ -221,7 +222,7 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): def memory_usage(self, deep: bool = False) -> int: ... @property def is_overlapping(self) -> bool: ... - def get_loc(self, key: Label) -> int | slice | npt.NDArray[np.bool_]: ... + def get_loc(self, key: Label) -> int | slice | np_1darray[np.bool]: ... @final def get_indexer( self, @@ -260,7 +261,7 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): @overload # type: ignore[override] def __gt__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_ndarray_bool: ... + ) -> np_1darray[np.bool]: ... @overload def __gt__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -268,7 +269,7 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): @overload # type: ignore[override] def __ge__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_ndarray_bool: ... + ) -> np_1darray[np.bool]: ... @overload def __ge__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -276,7 +277,7 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): @overload # type: ignore[override] def __le__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_ndarray_bool: ... + ) -> np_1darray[np.bool]: ... @overload def __le__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -284,13 +285,13 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): @overload # type: ignore[override] def __lt__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_ndarray_bool: ... + ) -> np_1darray[np.bool]: ... @overload def __lt__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] ) -> pd.Series[bool]: ... @overload # type: ignore[override] - def __eq__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + def __eq__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __eq__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[overload-overlap] @overload @@ -298,7 +299,7 @@ class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): self, other: object ) -> Literal[False]: ... @overload # type: ignore[override] - def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __ne__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[overload-overlap] @overload diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index d51528d9d..a4301ad29 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -25,8 +25,8 @@ from pandas._typing import ( MaskType, NaPosition, SequenceNotStr, + np_1darray, np_ndarray_anyint, - np_ndarray_bool, ) class MultiIndex(Index): @@ -161,4 +161,4 @@ class MultiIndex(Index): def equal_levels(self, other): ... def insert(self, loc, item): ... def delete(self, loc): ... - def isin(self, values, level=...) -> np_ndarray_bool: ... + def isin(self, values, level=...) -> np_1darray[np.bool]: ... diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 3dafef991..da025b8a4 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/pandas-stubs/core/indexes/period.pyi @@ -5,6 +5,7 @@ from typing import ( overload, ) +import numpy as np import pandas as pd from pandas import Index from pandas.core.indexes.accessors import PeriodIndexFieldOps @@ -19,7 +20,7 @@ from pandas._libs.tslibs import ( ) from pandas._libs.tslibs.period import _PeriodAddSub -class PeriodIndex(DatetimeIndexOpsMixin[pd.Period], PeriodIndexFieldOps): +class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexFieldOps): def __new__( cls, data=..., diff --git a/pandas-stubs/core/indexes/range.pyi b/pandas-stubs/core/indexes/range.pyi index 5c2d18263..1494b668c 100644 --- a/pandas-stubs/core/indexes/range.pyi +++ b/pandas-stubs/core/indexes/range.pyi @@ -8,16 +8,19 @@ from typing import ( ) import numpy as np -from pandas.core.indexes.base import Index +from pandas.core.indexes.base import ( + Index, + _IndexSubclassBase, +) from pandas._typing import ( HashableT, MaskType, + np_1darray, np_ndarray_anyint, - npt, ) -class RangeIndex(Index[int]): +class RangeIndex(_IndexSubclassBase[int, np.int64]): def __new__( cls, start: int | RangeIndex | range = ..., @@ -58,7 +61,7 @@ class RangeIndex(Index[int]): def argsort(self, *args, **kwargs): ... def factorize( self, sort: bool = False, use_na_sentinel: bool = True - ) -> tuple[npt.NDArray[np.intp], RangeIndex]: ... + ) -> tuple[np_1darray[np.intp], RangeIndex]: ... def equals(self, other): ... @final def join( diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 4c16f6e5d..28613dee3 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -33,7 +33,9 @@ from pandas._typing import ( num, ) -class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties): +class TimedeltaIndex( + DatetimeTimedeltaMixin[Timedelta, np.timedelta64], TimedeltaIndexProperties +): def __new__( cls, data: ( diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 69dfc1b08..f0f59ffc0 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -29,6 +29,7 @@ from typing import ( TypeVar, final, overload, + type_check_only, ) from matplotlib.axes import ( @@ -126,11 +127,14 @@ from pandas._typing import ( CompressionOptions, DropKeep, Dtype, + DTypeLike, DtypeObj, FilePath, FillnaOptions, FloatDtypeArg, FloatFormatType, + GenericT, + GenericT_co, GroupByObjectNonScalar, HashableT1, IgnoreRaise, @@ -165,6 +169,7 @@ from pandas._typing import ( StrDtypeArg, StrLike, Suffixes, + SupportsDType, T as _T, TimeAmbiguous, TimedeltaDtypeArg, @@ -176,6 +181,7 @@ from pandas._typing import ( ValueKeyFunc, VoidDtypeArg, WriteBuffer, + np_1darray, np_ndarray_anyint, np_ndarray_bool, np_ndarray_complex, @@ -453,7 +459,7 @@ class Series(IndexOpsMixin[S1], NDFrame): ): ... def __array__( self, dtype: _str | np.dtype = ..., copy: bool | None = ... - ) -> np.ndarray: ... + ) -> np_1darray: ... @property def axes(self) -> list: ... @final @@ -2878,14 +2884,6 @@ class Series(IndexOpsMixin[S1], NDFrame): **kwargs: Any, ) -> S1: ... def to_list(self) -> list[S1]: ... - @final - def to_numpy( - self, - dtype: npt.DTypeLike | None = ..., - copy: bool = False, - na_value: Scalar = ..., - **kwargs: Any, - ) -> np.ndarray: ... def tolist(self) -> list[S1]: ... def var( self, @@ -2946,7 +2944,34 @@ class Series(IndexOpsMixin[S1], NDFrame): @final def __bool__(self) -> NoReturn: ... -class TimestampSeries(Series[Timestamp]): +@type_check_only +class _SeriesSubclassBase(Series[S1], Generic[S1, GenericT_co]): + @overload + def to_numpy( # pyrefly: ignore + self, + dtype: None = None, + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT_co]: ... + @overload + def to_numpy( + self, + dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT], + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray[GenericT]: ... + @overload + def to_numpy( + self, + dtype: DTypeLike, + copy: bool = False, + na_value: Scalar = ..., + **kwargs, + ) -> np_1darray: ... + +class TimestampSeries(_SeriesSubclassBase[Timestamp, np.datetime64]): @property def dt(self) -> TimestampProperties: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] def __add__(self, other: TimedeltaSeries | np.timedelta64 | timedelta | BaseOffset) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] @@ -2998,7 +3023,7 @@ class TimestampSeries(Series[Timestamp]): **kwargs: Any, ) -> Never: ... -class TimedeltaSeries(Series[Timedelta]): +class TimedeltaSeries(_SeriesSubclassBase[Timedelta, np.timedelta64]): # ignores needed because of mypy @overload # type: ignore[override] def __add__(self, other: Period) -> PeriodSeries: ... @@ -3110,7 +3135,7 @@ class TimedeltaSeries(Series[Timedelta]): **kwargs: Any, ) -> Never: ... -class PeriodSeries(Series[Period]): +class PeriodSeries(_SeriesSubclassBase[Period, np.object_]): @property def dt(self) -> PeriodProperties: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] def __sub__(self, other: PeriodSeries) -> OffsetSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] @@ -3123,7 +3148,7 @@ class PeriodSeries(Series[Period]): **kwargs: Any, ) -> Never: ... -class OffsetSeries(Series[BaseOffset]): +class OffsetSeries(_SeriesSubclassBase[BaseOffset, np.object_]): @overload # type: ignore[override] def __radd__(self, other: Period) -> PeriodSeries: ... @overload @@ -3138,7 +3163,9 @@ class OffsetSeries(Series[BaseOffset]): **kwargs: Any, ) -> Never: ... -class IntervalSeries(Series[Interval[_OrderableT]], Generic[_OrderableT]): +class IntervalSeries( + _SeriesSubclassBase[Interval[_OrderableT], np.object_], Generic[_OrderableT] +): @property def array(self) -> IntervalArray: ... def diff(self, periods: int = ...) -> Never: ... diff --git a/pandas-stubs/core/strings/accessor.pyi b/pandas-stubs/core/strings/accessor.pyi index bedd8ac7b..a290fa12f 100644 --- a/pandas-stubs/core/strings/accessor.pyi +++ b/pandas-stubs/core/strings/accessor.pyi @@ -31,7 +31,7 @@ from pandas._typing import ( DtypeObj, Scalar, T, - np_ndarray_bool, + np_1darray, ) # Used for the result of str.split with expand=True @@ -39,7 +39,7 @@ _T_EXPANDING = TypeVar("_T_EXPANDING", bound=DataFrame | MultiIndex) # Used for the result of str.split with expand=False _T_LIST_STR = TypeVar("_T_LIST_STR", bound=Series[list[str]] | Index[list[str]]) # Used for the result of str.match -_T_BOOL = TypeVar("_T_BOOL", bound=Series[bool] | np_ndarray_bool) +_T_BOOL = TypeVar("_T_BOOL", bound=Series[bool] | np_1darray[np.bool]) # Used for the result of str.index / str.find _T_INT = TypeVar("_T_INT", bound=Series[int] | Index[int]) # Used for the result of str.encode diff --git a/tests/__init__.py b/tests/__init__.py index c4a94e773..89dd623e2 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -11,6 +11,10 @@ TYPE_CHECKING, Final, Literal, + TypeAlias, + TypeVar, + get_args, + get_origin, ) import numpy as np @@ -40,13 +44,21 @@ TimestampDtypeArg as TimestampDtypeArg, UIntDtypeArg as UIntDtypeArg, VoidDtypeArg as VoidDtypeArg, + np_1darray as np_1darray, + np_2darray as np_2darray, + np_ndarray as np_ndarray, np_ndarray_bool as np_ndarray_bool, np_ndarray_int as np_ndarray_int, ) else: + _G = TypeVar("_G", bound=np.generic) + _S = TypeVar("_S", bound=tuple[int, ...]) # Separately define here so pytest works - np_ndarray_bool = npt.NDArray[np.bool_] - np_ndarray_int = npt.NDArray[np.signedinteger] + np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]] + np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]] + np_ndarray: TypeAlias = np.ndarray[_S, np.dtype[_G]] + np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_] + np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger] TYPE_CHECKING_INVALID_USAGE: Final = TYPE_CHECKING WINDOWS = os.name == "nt" or "cygwin" in platform.system().lower() @@ -61,8 +73,39 @@ def check( attr: str = "left", index_to_check_for_type: Literal[0, -1] = 0, ) -> T: - if not isinstance(actual, klass): + __tracebackhide__ = True + origin = get_origin(klass) + if not isinstance(actual, origin or klass): raise RuntimeError(f"Expected type '{klass}' but got '{type(actual)}'") + if origin is np.ndarray: + # Check shape and dtype + args = get_args(klass) + shape_type = args[0] if len(args) >= 1 else None + dtype_type = args[1] if len(args) >= 2 else None + if ( + shape_type + and get_origin(shape_type) is tuple + and (tuple_args := get_args(shape_type)) + and ... not in tuple_args # fixed-length tuple + and (arr_ndim := getattr(actual, "ndim")) + != (expected_ndim := len(tuple_args)) + ): + raise RuntimeError( + f"Array has wrong dimension {arr_ndim}, expected {expected_ndim}" + ) + + if ( + dtype_type + and get_origin(dtype_type) is np.dtype + and (dtype_args := get_args(dtype_type)) + and isinstance((expected_dtype := dtype_args[0]), type) + and issubclass(expected_dtype, np.generic) + and (arr_dtype := getattr(actual, "dtype")) != expected_dtype + ): + raise RuntimeError( + f"Array has wrong dtype {arr_dtype}, expected {expected_dtype.__name__}" + ) + if dtype is None: return actual diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 795193782..4070b4b3d 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -56,6 +56,7 @@ WINDOWS, check, ensure_clean, + np_1darray, pytest_warns_bounded, ) from tests.extension.decimal.array import DecimalDtype @@ -91,7 +92,6 @@ UIntDtypeArg, VoidDtypeArg, ) - from tests import np_ndarray_int # noqa: F401 else: TimedeltaSeries: TypeAlias = pd.Series @@ -1870,7 +1870,7 @@ def test_types_to_dict() -> None: def test_categorical_codes(): # GH-111 cat = pd.Categorical(["a", "b", "a"]) - assert_type(cat.codes, "np_ndarray_int") + check(assert_type(cat.codes, np_1darray[np.signedinteger]), np_1darray[np.int8]) def test_relops() -> None: @@ -2022,12 +2022,12 @@ def test_dtype_type() -> None: def test_types_to_numpy() -> None: s = pd.Series(["a", "b", "c"], dtype=str) - check(assert_type(s.to_numpy(), np.ndarray), np.ndarray) - check(assert_type(s.to_numpy(dtype="str", copy=True), np.ndarray), np.ndarray) - check(assert_type(s.to_numpy(na_value=0), np.ndarray), np.ndarray) - check(assert_type(s.to_numpy(na_value=np.int32(4)), np.ndarray), np.ndarray) - check(assert_type(s.to_numpy(na_value=np.float16(4)), np.ndarray), np.ndarray) - check(assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np.ndarray), np.ndarray) + check(assert_type(s.to_numpy(), np_1darray), np_1darray) + check(assert_type(s.to_numpy(dtype="str", copy=True), np_1darray), np_1darray) + check(assert_type(s.to_numpy(na_value=0), np_1darray), np_1darray) + check(assert_type(s.to_numpy(na_value=np.int32(4)), np_1darray), np_1darray) + check(assert_type(s.to_numpy(na_value=np.float16(4)), np_1darray), np_1darray) + check(assert_type(s.to_numpy(na_value=np.complex128(4, 7)), np_1darray), np_1darray) def test_where() -> None: diff --git a/tests/test_frame.py b/tests/test_frame.py index f70504ef4..cf1b58015 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -57,6 +57,7 @@ TYPE_CHECKING_INVALID_USAGE, check, ensure_clean, + np_2darray, pytest_warns_bounded, ) @@ -1946,18 +1947,18 @@ def test_types_cov() -> None: def test_types_to_numpy() -> None: df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5]}) - check(assert_type(df.to_numpy(), np.ndarray), np.ndarray) - check(assert_type(df.to_numpy(dtype="str", copy=True), np.ndarray), np.ndarray) + check(assert_type(df.to_numpy(), np_2darray), np_2darray) + check(assert_type(df.to_numpy(dtype="str", copy=True), np_2darray), np_2darray) # na_value param was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html - check(assert_type(df.to_numpy(na_value=0), np.ndarray), np.ndarray) + check(assert_type(df.to_numpy(na_value=0), np_2darray), np_2darray) df = pd.DataFrame(data={"col1": [1, 1, 2]}, dtype=np.complex128) - check(assert_type(df.to_numpy(na_value=0), np.ndarray), np.ndarray) - check(assert_type(df.to_numpy(na_value=np.int32(4)), np.ndarray), np.ndarray) - check(assert_type(df.to_numpy(na_value=np.float16(3.68)), np.ndarray), np.ndarray) + check(assert_type(df.to_numpy(na_value=0), np_2darray), np_2darray) + check(assert_type(df.to_numpy(na_value=np.int32(4)), np_2darray), np_2darray) + check(assert_type(df.to_numpy(na_value=np.float16(3.68)), np_2darray), np_2darray) check( - assert_type(df.to_numpy(na_value=np.complex128(3.8, -493.2)), np.ndarray), - np.ndarray, + assert_type(df.to_numpy(na_value=np.complex128(3.8, -493.2)), np_2darray), + np_2darray, ) diff --git a/tests/test_indexes.py b/tests/test_indexes.py index f6ad2a270..819e5b248 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -26,6 +26,7 @@ PD_LTE_23, TYPE_CHECKING_INVALID_USAGE, check, + np_1darray, pytest_warns_bounded, ) @@ -42,13 +43,13 @@ def test_index_duplicated() -> None: df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2])) ind = df.index duplicated = ind.duplicated("first") - check(assert_type(duplicated, npt.NDArray[np.bool_]), np.ndarray, np.bool_) + check(assert_type(duplicated, np_1darray[np.bool]), np_1darray[np.bool]) def test_index_isin() -> None: ind = pd.Index([1, 2, 3, 4, 5]) isin = ind.isin([2, 4]) - check(assert_type(isin, npt.NDArray[np.bool_]), np.ndarray, np.bool_) + check(assert_type(isin, np_1darray[np.bool]), np_1darray[np.bool]) def test_index_astype() -> None: @@ -202,13 +203,6 @@ def test_str_rsplit() -> None: ) -def test_str_match() -> None: - i = pd.Index( - ["applep", "bananap", "Cherryp", "DATEp", "eGGpLANTp", "123p", "23.45p"] - ) - check(assert_type(i.str.match("pp"), npt.NDArray[np.bool_]), np.ndarray, np.bool_) - - def test_index_rename() -> None: """Test that index rename returns an element of type Index.""" ind = pd.Index([1, 2, 3], name="foo") @@ -244,9 +238,25 @@ def test_index_neg(): def test_types_to_numpy() -> None: idx = pd.Index([1, 2]) - check(assert_type(idx.to_numpy(), np.ndarray), np.ndarray) - check(assert_type(idx.to_numpy(dtype="int", copy=True), np.ndarray), np.ndarray) - check(assert_type(idx.to_numpy(na_value=0), np.ndarray), np.ndarray) + check(assert_type(idx.to_numpy(), np_1darray), np_1darray) + check(assert_type(idx.to_numpy(dtype="int", copy=True), np_1darray), np_1darray) + check(assert_type(idx.to_numpy(na_value=0), np_1darray), np_1darray) + + r_idx = pd.RangeIndex(2) + check(assert_type(r_idx.to_numpy(), np_1darray[np.int64]), np_1darray[np.int64]) + check( + assert_type(r_idx.to_numpy(na_value=0), np_1darray[np.int64]), + np_1darray[np.int64], + ) + check( + assert_type(r_idx.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(r_idx.to_numpy(dtype=np.int32), np_1darray[np.int32]), + np_1darray[np.int32], + ) def test_index_arithmetic() -> None: @@ -289,10 +299,10 @@ def test_index_relops() -> None: check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex) ind = pd.Index([1, 2, 3]) - check(assert_type(ind <= 2, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(ind >= 2, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(ind < 2, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(ind > 2, npt.NDArray[np.bool_]), np.ndarray, np.bool_) + check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool]) def test_range_index_union(): @@ -1299,16 +1309,14 @@ def test_datetime_operators_builtin() -> None: def test_get_loc() -> None: unique_index = pd.Index(list("abc")) check( - assert_type( - unique_index.get_loc("b"), Union[int, slice, npt.NDArray[np.bool_]] - ), + assert_type(unique_index.get_loc("b"), Union[int, slice, np_1darray[np.bool]]), int, ) monotonic_index = pd.Index(list("abbc")) check( assert_type( - monotonic_index.get_loc("b"), Union[int, slice, npt.NDArray[np.bool_]] + monotonic_index.get_loc("b"), Union[int, slice, np_1darray[np.bool]] ), slice, ) @@ -1316,10 +1324,25 @@ def test_get_loc() -> None: non_monotonic_index = pd.Index(list("abcb")) check( assert_type( - non_monotonic_index.get_loc("b"), Union[int, slice, npt.NDArray[np.bool_]] + non_monotonic_index.get_loc("b"), Union[int, slice, np_1darray[np.bool]] + ), + np_1darray[np.bool], + ) + + i1, i2, i3 = pd.Interval(0, 1), pd.Interval(1, 2), pd.Interval(0, 2) + unique_interval_index = pd.IntervalIndex([i1, i2]) + check( + assert_type( + unique_interval_index.get_loc(i1), Union[int, slice, np_1darray[np.bool]] + ), + np.int64, + ) + overlap_interval_index = pd.IntervalIndex([i1, i2, i3]) + check( + assert_type( + overlap_interval_index.get_loc(1), Union[int, slice, np_1darray[np.bool]] ), - np.ndarray, - np.bool_, + np_1darray[np.bool], ) @@ -1336,14 +1359,14 @@ def test_value_counts() -> None: def test_index_factorize() -> None: """Test Index.factorize method.""" codes, idx_uniques = pd.Index(["b", "b", "a", "c", "b"]).factorize() - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(idx_uniques, np.ndarray | Index | Categorical), pd.Index) + check(assert_type(codes, np_1darray), np_1darray) + check(assert_type(idx_uniques, np_1darray | Index | Categorical), pd.Index) codes, idx_uniques = pd.Index(["b", "b", "a", "c", "b"]).factorize( use_na_sentinel=False ) - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(idx_uniques, np.ndarray | Index | Categorical), pd.Index) + check(assert_type(codes, np_1darray), np_1darray) + check(assert_type(idx_uniques, np_1darray | Index | Categorical), pd.Index) def test_disallow_empty_index() -> None: diff --git a/tests/test_interval.py b/tests/test_interval.py index 208255f89..8d9081710 100644 --- a/tests/test_interval.py +++ b/tests/test_interval.py @@ -1,13 +1,13 @@ from __future__ import annotations import numpy as np -from numpy import typing as npt import pandas as pd from typing_extensions import assert_type from tests import ( TYPE_CHECKING_INVALID_USAGE, check, + np_1darray, ) @@ -127,4 +127,4 @@ def test_interval_array_contains(): ser = pd.Series(obj, index=df.index) arr = ser.array check(assert_type(arr.contains(df["A"]), "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(arr.contains(3), npt.NDArray[np.bool_]), np.ndarray) + check(assert_type(arr.contains(3), np_1darray[np.bool]), np_1darray[np.bool]) diff --git a/tests/test_pandas.py b/tests/test_pandas.py index cfa6f6be2..a7633fa8a 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -33,6 +33,9 @@ PD_LTE_23, TYPE_CHECKING_INVALID_USAGE, check, + np_1darray, + np_2darray, + np_ndarray_bool, pytest_warns_bounded, ) @@ -428,20 +431,47 @@ def test_types_json_normalize() -> None: def test_isna() -> None: # https://github.com/pandas-dev/pandas-stubs/issues/264 - s1 = pd.Series([1, np.nan, 3.2]) - check(assert_type(pd.isna(s1), "pd.Series[bool]"), pd.Series, np.bool_) - - s2 = pd.Series([1, 3.2]) - check(assert_type(pd.notna(s2), "pd.Series[bool]"), pd.Series, np.bool_) - - df1 = pd.DataFrame({"a": [1, 2, 1, 2], "b": [1, 1, 2, np.nan]}) - check(assert_type(pd.isna(df1), "pd.DataFrame"), pd.DataFrame) - - idx1 = pd.Index([1, 2, np.nan]) - check(assert_type(pd.isna(idx1), npt.NDArray[np.bool_]), np.ndarray, np.bool_) - - idx2 = pd.Index([1, 2]) - check(assert_type(pd.notna(idx2), npt.NDArray[np.bool_]), np.ndarray, np.bool_) + s = pd.Series([1, np.nan, 3.2]) + check(assert_type(pd.isna(s), "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(pd.notna(s), "pd.Series[bool]"), pd.Series, np.bool_) + + df = pd.DataFrame({"a": [1, 2, 1, 2], "b": [1, 1, 2, np.nan]}) + check(assert_type(pd.isna(df), "pd.DataFrame"), pd.DataFrame) + check(assert_type(pd.notna(df), "pd.DataFrame"), pd.DataFrame) + + idx = pd.Index([1, 2, np.nan, float("nan")]) + check(assert_type(pd.isna(idx), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(pd.notna(idx), np_1darray[np.bool]), np_1darray[np.bool]) + + # ExtensionArray + ext_arr = idx.array + check(assert_type(pd.isna(ext_arr), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(pd.notna(ext_arr), np_1darray[np.bool]), np_1darray[np.bool]) + + # 1-D numpy array + arr_1d = idx.to_numpy() + check(assert_type(pd.isna(arr_1d), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(pd.notna(arr_1d), np_1darray[np.bool]), np_1darray[np.bool]) + + # 2-D numpy array + arr_2d = idx.to_numpy().reshape(2, 2) + check(assert_type(pd.isna(arr_2d), np_2darray[np.bool]), np_2darray[np.bool]) + check(assert_type(pd.notna(arr_2d), np_2darray[np.bool]), np_2darray[np.bool]) + + # N-D numpy array + arr_nd = idx.to_numpy().reshape([2, 2]) + check(assert_type(pd.isna(arr_nd), np_ndarray_bool), np_ndarray_bool) + check(assert_type(pd.notna(arr_nd), np_ndarray_bool), np_ndarray_bool) + + # List of scalars + l_sca = [1, 2.5, float("nan")] + check(assert_type(pd.isna(l_sca), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(pd.notna(l_sca), np_1darray[np.bool]), np_1darray[np.bool]) + + # List of unknown members + l_any: list[object] = [arr_1d, ext_arr] + check(assert_type(pd.isna(l_any), np_ndarray_bool), np_ndarray_bool) + check(assert_type(pd.notna(l_any), np_ndarray_bool), np_ndarray_bool) assert check(assert_type(pd.isna(pd.NA), bool), bool) assert not check(assert_type(pd.notna(pd.NA), bool), bool) @@ -909,15 +939,15 @@ def test_factorize() -> None: check(assert_type(uniques, np.ndarray), np.ndarray) codes, cat_uniques = pd.factorize(pd.Categorical(["b", "b", "a", "c", "b"])) - check(assert_type(codes, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray), np_1darray) check(assert_type(cat_uniques, pd.Categorical), pd.Categorical) codes, idx_uniques = pd.factorize(pd.Index(["b", "b", "a", "c", "b"])) - check(assert_type(codes, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray), np_1darray) check(assert_type(idx_uniques, pd.Index), pd.Index) codes, idx_uniques = pd.factorize(pd.Series(["b", "b", "a", "c", "b"])) - check(assert_type(codes, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray), np_1darray) check(assert_type(idx_uniques, pd.Index), pd.Index) codes, uniques = pd.factorize(np.array(list("bbacb"))) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index bfc78ae65..92b29f4bc 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -26,6 +26,9 @@ from tests import ( TYPE_CHECKING_INVALID_USAGE, check, + np_1darray, + np_2darray, + np_ndarray_bool, pytest_warns_bounded, ) @@ -48,8 +51,6 @@ PeriodSeries: TypeAlias = pd.Series OffsetSeries: TypeAlias = pd.Series -from tests import np_ndarray_bool - def test_interval() -> None: interval_i = pd.Interval(0, 1, closed="left") @@ -314,73 +315,55 @@ def test_interval_cmp(): interval_index_int = pd.IntervalIndex([interval_i]) check( - assert_type(interval_index_int >= interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int >= interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_index_int < interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int < interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_index_int <= interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int <= interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_index_int > interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int > interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_i >= interval_index_int, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_i >= interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_i < interval_index_int, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_i < interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_i <= interval_index_int, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_i <= interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_i > interval_index_int, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_i > interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_index_int == interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int == interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type(interval_index_int != interval_i, np_ndarray_bool), - np.ndarray, - np.bool_, + assert_type(interval_index_int != interval_i, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type( - interval_i == interval_index_int, - np_ndarray_bool, - ), - np.ndarray, - np.bool_, + assert_type(interval_i == interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) check( - assert_type( - interval_i != interval_index_int, - np_ndarray_bool, - ), - np.ndarray, - np.bool_, + assert_type(interval_i != interval_index_int, np_1darray[np.bool]), + np_1darray[np.bool], ) @@ -838,248 +821,201 @@ def test_timedelta_mod_abs_unary() -> None: check(assert_type(+td, pd.Timedelta), pd.Timedelta) -def test_timedelta_cmp() -> None: +def test_timedelta_cmp_scalar() -> None: td = pd.Timedelta("1 day") - ndarray_td64: npt.NDArray[np.timedelta64] = np.array( - [1, 2, 3], dtype="timedelta64[D]" - ) - c_timedelta = td - c_dt_timedelta = dt.timedelta(days=1) - c_timedelta64 = np.timedelta64(1, "D") - c_ndarray_td64 = ndarray_td64 - c_timedelta_index = pd.to_timedelta([1, 2, 3], unit="D") - c_timedelta_series = pd.Series(pd.TimedeltaIndex([1, 2, 3])) - - check(assert_type(td < c_timedelta, bool), bool) - check(assert_type(td < c_dt_timedelta, bool), bool) - check(assert_type(td < c_timedelta64, bool), bool) - check(assert_type(td < c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_dt_timedelta < td, bool), bool) - check(assert_type(c_ndarray_td64 < td, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_) - - gt = check(assert_type(td > c_timedelta, bool), bool) - le = check(assert_type(td <= c_timedelta, bool), bool) - assert gt != le - - gt = check(assert_type(td > c_dt_timedelta, bool), bool) - le = check(assert_type(td <= c_dt_timedelta, bool), bool) - assert gt != le - - gt = check(assert_type(td > c_timedelta64, bool), bool) - le = check(assert_type(td <= c_timedelta64, bool), bool) - assert gt != le - - gt_a = check( - assert_type(td > c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - le_a = check( - assert_type(td <= c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - gt_a = check( - assert_type(td > c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - le_a = check( - assert_type(td <= c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - gt_s = check( - assert_type(td > c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - le_s = check( - assert_type(td <= c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (gt_s != le_s).all() - - gt = check(assert_type(c_dt_timedelta > td, bool), bool) - le = check(assert_type(c_dt_timedelta <= td, bool), bool) - assert gt != le - - gt_b = check(assert_type(c_timedelta64 > td, np.bool), bool) - le_b = check(assert_type(c_timedelta64 <= td, np.bool), bool) - assert gt_b != le_b - - gt_a = check( - assert_type(c_ndarray_td64 > td, np_ndarray_bool), np.ndarray, np.bool_ - ) - le_a = check( - assert_type(c_ndarray_td64 <= td, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - gt_a = check( - assert_type(c_timedelta_index > td, np_ndarray_bool), np.ndarray, np.bool_ - ) - le_a = check( - assert_type(c_timedelta_index <= td, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - eq_s = check( - assert_type(c_timedelta_series > td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(c_timedelta_series <= td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - lt = check(assert_type(td < c_timedelta, bool), bool) - ge = check(assert_type(td >= c_timedelta, bool), bool) - assert lt != ge - - lt = check(assert_type(td < c_dt_timedelta, bool), bool) - ge = check(assert_type(td >= c_dt_timedelta, bool), bool) - assert lt != ge - - lt = check(assert_type(td < c_timedelta64, bool), bool) - ge = check(assert_type(td >= c_timedelta64, bool), bool) - assert lt != ge - - lt_a = check( - assert_type(td < c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - ge_a = check( - assert_type(td >= c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - lt_a = check( - assert_type(td < c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - ge_a = check( - assert_type(td >= c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - eq_s = check( - assert_type(td < c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(td >= c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - lt = check(assert_type(c_dt_timedelta < td, bool), bool) - ge = check(assert_type(c_dt_timedelta >= td, bool), bool) - assert lt != ge - - lt_b = check(assert_type(c_timedelta64 < td, np.bool), bool) - ge_b = check(assert_type(c_timedelta64 >= td, np.bool), bool) - assert lt_b != ge_b - - lt_a = check( - assert_type(c_ndarray_td64 < td, np_ndarray_bool), np.ndarray, np.bool_ - ) - ge_a = check( - assert_type(c_ndarray_td64 >= td, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - lt_a = check( - assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_ - ) - ge_a = check( - assert_type(c_timedelta_index >= td, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - eq_s = check( - assert_type(c_timedelta_series < td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(c_timedelta_series >= td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - eq = check(assert_type(td == td, bool), bool) - ne = check(assert_type(td != td, bool), bool) - assert eq != ne - - eq = check(assert_type(td == c_dt_timedelta, bool), bool) - ne = check(assert_type(td != c_dt_timedelta, bool), bool) - assert eq != ne - - eq = check(assert_type(td == c_timedelta64, bool), bool) - ne = check(assert_type(td != c_timedelta64, bool), bool) - assert eq != ne - - eq_a = check( - assert_type(td == c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_a = check( - assert_type(td != c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_a != ne_a).all() - - eq_a = check( - assert_type(td == c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_a = check( - assert_type(td != c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_a != ne_a).all() - - eq_s = check( - assert_type(td == c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(td != c_timedelta_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - eq = check(assert_type(td == 1, Literal[False]), bool) - ne = check(assert_type(td != 1, Literal[True]), bool) - assert eq != ne - - eq = check(assert_type(td == (3 + 2j), Literal[False]), bool) - ne = check(assert_type(td != (3 + 2j), Literal[True]), bool) - assert eq != ne - - -def test_timedelta_cmp_rhs() -> None: - # Test that check eq and ne when Timedelta is the RHS argument - # that use the __eq__ and __ne__ methods of the LHS + td2 = pd.Timedelta("1 hour") + py_td = dt.timedelta(days=1) + np_td = np.timedelta64(1, "D") + + # >, <= + gt1 = check(assert_type(td > td2, bool), bool) + le1 = check(assert_type(td <= td2, bool), bool) + assert gt1 != le1 + gt2 = check(assert_type(td > py_td, bool), bool) + le2 = check(assert_type(td <= py_td, bool), bool) + assert gt2 != le2 + gt3 = check(assert_type(td > np_td, bool), bool) + le3 = check(assert_type(td <= np_td, bool), bool) + assert gt3 != le3 + gt4 = check(assert_type(py_td > td, bool), bool) + le4 = check(assert_type(py_td <= td, bool), bool) + assert gt4 != le4 + gt5 = check(assert_type(np_td > td, np.bool), bool) + le5 = check(assert_type(np_td <= td, np.bool), bool) + assert gt5 != le5 + + # <, >= + lt1 = check(assert_type(td < td2, bool), bool) + ge1 = check(assert_type(td >= td2, bool), bool) + assert lt1 != ge1 + lt2 = check(assert_type(td < py_td, bool), bool) + ge2 = check(assert_type(td >= py_td, bool), bool) + assert lt2 != ge2 + lt3 = check(assert_type(td < np_td, bool), bool) + ge3 = check(assert_type(td >= np_td, bool), bool) + assert lt3 != ge3 + lt4 = check(assert_type(py_td < td, bool), bool) + ge4 = check(assert_type(py_td >= td, bool), bool) + assert lt4 != ge4 + lt5 = check(assert_type(np_td < td, np.bool), bool) + ge5 = check(assert_type(np_td >= td, np.bool), bool) + assert lt5 != ge5 + + # ==, != + eq1 = check(assert_type(td == td, bool), bool) + ne1 = check(assert_type(td != td, bool), bool) + assert eq1 != ne1 + eq2 = check(assert_type(td == py_td, bool), bool) + ne2 = check(assert_type(td != py_td, bool), bool) + assert eq2 != ne2 + eq3 = check(assert_type(td == np_td, bool), bool) + ne3 = check(assert_type(td != np_td, bool), bool) + assert eq3 != ne3 + eq4 = check(assert_type(td == 1, Literal[False]), bool) + ne4 = check(assert_type(td != 1, Literal[True]), bool) + assert eq4 != ne4 + eq5 = check(assert_type(td == (3 + 2j), Literal[False]), bool) + ne5 = check(assert_type(td != (3 + 2j), Literal[True]), bool) + assert eq5 != ne5 + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(py_td == td, bool), bool) + ne_rhs1 = check(assert_type(py_td != td, bool), bool) + assert eq_rhs1 != ne_rhs1 + eq_rhs2 = check(assert_type(np_td == td, Any), bool) + ne_rhs2 = check(assert_type(np_td != td, Any), bool) + assert eq_rhs2 != ne_rhs2 + + +def test_timedelta_cmp_series() -> None: td = pd.Timedelta("1 day") - ndarray_td64: npt.NDArray[np.timedelta64] = np.array( - [1, 2, 3], dtype="timedelta64[D]" - ) - c_dt_timedelta = dt.timedelta(days=1) - c_timedelta64 = np.timedelta64(1, "D") - c_ndarray_td64 = ndarray_td64 - c_timedelta_index = pd.to_timedelta([1, 2, 3], unit="D") - c_timedelta_series = pd.Series(pd.TimedeltaIndex([1, 2, 3])) - - eq = check(assert_type(c_dt_timedelta == td, bool), bool) - ne = check(assert_type(c_dt_timedelta != td, bool), bool) - assert eq != ne - - eq = check(assert_type(c_timedelta64 == td, Any), bool) - ne = check(assert_type(c_timedelta64 != td, Any), bool) - assert eq != ne - - eq_a = check(assert_type(c_ndarray_td64 == td, Any), np.ndarray, np.bool_) - ne_a = check(assert_type(c_ndarray_td64 != td, Any), np.ndarray, np.bool_) - assert (eq_a != ne_a).all() - - eq_a = check( - assert_type(c_timedelta_index == td, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_a = check( - assert_type(c_timedelta_index != td, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_a != ne_a).all() - - eq_s = check( - assert_type(c_timedelta_series == td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(c_timedelta_series != td, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() + td_ser = pd.Series(pd.TimedeltaIndex([1, 2, 3])) # TimedeltaSeries + + # >, <= + gt1 = check(assert_type(td > td_ser, "pd.Series[bool]"), pd.Series, np.bool) + le1 = check(assert_type(td <= td_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt1 != le1).all() + gt2 = check(assert_type(td_ser > td, "pd.Series[bool]"), pd.Series, np.bool) + le2 = check(assert_type(td_ser <= td, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt2 != le2).all() + + # <, >= + lt1 = check(assert_type(td < td_ser, "pd.Series[bool]"), pd.Series, np.bool) + ge1 = check(assert_type(td >= td_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt1 != ge1).all() + lt2 = check(assert_type(td_ser < td, "pd.Series[bool]"), pd.Series, np.bool) + ge2 = check(assert_type(td_ser >= td, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt2 != ge2).all() + + # ==, != + eq1 = check(assert_type(td == td_ser, "pd.Series[bool]"), pd.Series, np.bool) + ne1 = check(assert_type(td != td_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq1 != ne1).all() + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(td_ser == td, "pd.Series[bool]"), pd.Series, np.bool) + ne_rhs1 = check(assert_type(td_ser != td, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq_rhs1 != ne_rhs1).all() + + +def test_timedelta_cmp_index() -> None: + td = pd.Timedelta("1 day") + td_idx = pd.to_timedelta([1, 2, 3], unit="D") # TimedeltaIndex + + # >, <= + gt1 = check(assert_type(td > td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + le1 = check(assert_type(td <= td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt1 != le1).all() + gt2 = check(assert_type(td_idx > td, np_1darray[np.bool]), np_1darray[np.bool]) + le2 = check(assert_type(td_idx <= td, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt2 != le2).all() + + # <, >= + lt1 = check(assert_type(td < td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ge1 = check(assert_type(td >= td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt1 != ge1).all() + lt2 = check(assert_type(td_idx < td, np_1darray[np.bool]), np_1darray[np.bool]) + ge2 = check(assert_type(td_idx >= td, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt2 != ge2).all() + + # ==, != + eq1 = check(assert_type(td == td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ne1 = check(assert_type(td != td_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq1 != ne1).all() + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(td_idx == td, np_1darray[np.bool]), np_1darray[np.bool]) + ne_rhs1 = check(assert_type(td_idx != td, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_rhs1 != ne_rhs1).all() + + +def test_timedelta_cmp_array() -> None: + td = pd.Timedelta("1 day") + arr_1d = pd.to_timedelta([1, 2, 3, 4], unit="D").to_numpy() + arr_2d = arr_1d.reshape(2, 2) + arr_nd = arr_1d.reshape([4]) + + # >, <= + gt_nd1 = check(assert_type(td > arr_nd, np_ndarray_bool), np.ndarray, np.bool) + le_nd1 = check(assert_type(td <= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd1 != le_nd1).all() + gt_nd2 = check(assert_type(arr_nd > td, np_ndarray_bool), np.ndarray, np.bool) + le_nd2 = check(assert_type(arr_nd <= td, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd2 != le_nd2).all() + gt_2d1 = check(assert_type(td > arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + le_2d1 = check(assert_type(td <= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (gt_2d1 != le_2d1).all() + gt_2d2 = check(assert_type(arr_2d > td, np_2darray[np.bool]), np_2darray[np.bool]) + le_2d2 = check(assert_type(arr_2d <= td, np_2darray[np.bool]), np_2darray[np.bool]) + assert (gt_2d2 != le_2d2).all() + gt_1d1 = check(assert_type(td > arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + le_1d1 = check(assert_type(td <= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_1d1 != le_1d1).all() + gt_1d2 = check(assert_type(arr_1d > td, np_1darray[np.bool]), np_1darray[np.bool]) + le_1d2 = check(assert_type(arr_1d <= td, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_1d2 != le_1d2).all() + + # <, >= + lt_nd1 = check(assert_type(td < arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ge_nd1 = check(assert_type(td >= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd1 != ge_nd1).all() + lt_nd2 = check(assert_type(arr_nd < td, np_ndarray_bool), np.ndarray, np.bool) + ge_nd2 = check(assert_type(arr_nd >= td, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd2 != ge_nd2).all() + lt_2d1 = check(assert_type(td < arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ge_2d1 = check(assert_type(td >= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (lt_2d1 != ge_2d1).all() + lt_2d2 = check(assert_type(arr_2d < td, np_2darray[np.bool]), np_2darray[np.bool]) + ge_2d2 = check(assert_type(arr_2d >= td, np_2darray[np.bool]), np_2darray[np.bool]) + assert (lt_2d2 != ge_2d2).all() + lt_1d1 = check(assert_type(td < arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ge_1d1 = check(assert_type(td >= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_1d1 != ge_1d1).all() + lt_1d2 = check(assert_type(arr_1d < td, np_1darray[np.bool]), np_1darray[np.bool]) + ge_1d2 = check(assert_type(arr_1d >= td, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_1d2 != ge_1d2).all() + + # ==, != + eq_nd1 = check(assert_type(td == arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ne_nd1 = check(assert_type(td != arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (eq_nd1 != ne_nd1).all() + eq_2d1 = check(assert_type(td == arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ne_2d1 = check(assert_type(td != arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (eq_2d1 != ne_2d1).all() + eq_1d1 = check(assert_type(td == arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ne_1d1 = check(assert_type(td != arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_1d1 != ne_1d1).all() + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs_nd1 = check(assert_type(arr_nd == td, Any), np_ndarray_bool) + ne_rhs_nd1 = check(assert_type(arr_nd != td, Any), np_ndarray_bool) + assert (eq_rhs_nd1 != ne_rhs_nd1).all() + eq_rhs_2d1 = check(assert_type(arr_2d == td, Any), np_2darray[np.bool]) + ne_rhs_2d1 = check(assert_type(arr_2d != td, Any), np_2darray[np.bool]) + assert (eq_rhs_2d1 != ne_rhs_2d1).all() + eq_rhs_1d1 = check(assert_type(arr_1d == td, Any), np_1darray[np.bool]) + ne_rhs_1d1 = check(assert_type(arr_1d != td, Any), np_1darray[np.bool]) + assert (eq_rhs_1d1 != ne_rhs_1d1).all() def test_timestamp_construction() -> None: @@ -1250,218 +1186,230 @@ def test_timestamp_add_sub() -> None: ) -def test_timestamp_cmp() -> None: +def test_timestamp_cmp_scalar() -> None: ts = pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27) - - np_dt64_arr: npt.NDArray[np.datetime64] = np.array( - [1, 2, 3], dtype="datetime64[ns]" - ) - - c_timestamp = ts - c_np_dt64 = np.datetime64(1, "ns") - c_dt_datetime = dt.datetime(year=2000, month=1, day=1) - c_datetimeindex = pd.DatetimeIndex(["2000-1-1"]) - # DatetimeIndex, but the type checker detects it to be UnknownIndex. - c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index - c_np_ndarray_dt64 = np_dt64_arr - c_series_dt64: TimestampSeries = pd.Series([1, 2, 3], dtype="datetime64[ns]") - c_series_timestamp = pd.Series(pd.DatetimeIndex(["2000-1-1"])) - check(assert_type(c_series_timestamp, TimestampSeries), pd.Series, pd.Timestamp) - # Use xor to ensure one is True and the other is False - # Correctness ensures since tested to be bools - gt = check(assert_type(ts > c_timestamp, bool), bool) - lte = check(assert_type(ts <= c_timestamp, bool), bool) - assert gt != lte - - gt = check(assert_type(ts > c_np_dt64, bool), bool) - lte = check(assert_type(ts <= c_np_dt64, bool), bool) - assert gt != lte - - gt = check(assert_type(ts > c_dt_datetime, bool), bool) - lte = check(assert_type(ts <= c_dt_datetime, bool), bool) - assert gt != lte - - check(assert_type(ts > c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts <= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts > c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts <= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(ts > c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts <= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(ts > c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(ts <= c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_) - - check(assert_type(ts > c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(ts <= c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_) - - check(assert_type(c_np_dt64 > ts, np.bool), bool) - check(assert_type(c_np_dt64 <= ts, np.bool), bool) - - gt = check(assert_type(c_dt_datetime > ts, bool), bool) - lte = check(assert_type(c_dt_datetime <= ts, bool), bool) - assert gt != lte - - check(assert_type(c_datetimeindex > ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_datetimeindex <= ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_unknown_index > ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_unknown_index <= ts, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(c_np_ndarray_dt64 > ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_np_ndarray_dt64 <= ts, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(c_series_dt64 > ts, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(c_series_dt64 <= ts, "pd.Series[bool]"), pd.Series, np.bool_) - - gte = check(assert_type(ts >= c_timestamp, bool), bool) - lt = check(assert_type(ts < c_timestamp, bool), bool) - assert gte != lt - - gte = check(assert_type(ts >= c_np_dt64, bool), bool) - lt = check(assert_type(ts < c_np_dt64, bool), bool) - assert gte != lt - - gte = check(assert_type(ts >= c_dt_datetime, bool), bool) - lt = check(assert_type(ts < c_dt_datetime, bool), bool) - assert gte != lt - - check(assert_type(ts >= c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts < c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts >= c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts < c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(ts >= c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(ts < c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(ts >= c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(ts < c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_) - - check(assert_type(ts >= c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(ts < c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_) - - gte = check(assert_type(c_dt_datetime >= ts, bool), bool) - lt = check(assert_type(c_dt_datetime < ts, bool), bool) - assert gte != lt - - check(assert_type(c_np_dt64 >= ts, np.bool), bool) - check(assert_type(c_np_dt64 < ts, np.bool), bool) - - check(assert_type(c_datetimeindex >= ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_datetimeindex < ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_unknown_index >= ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_unknown_index < ts, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(c_np_ndarray_dt64 >= ts, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c_np_ndarray_dt64 < ts, np_ndarray_bool), np.ndarray, np.bool_) - - check(assert_type(c_series_dt64 >= ts, "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(c_series_dt64 < ts, "pd.Series[bool]"), pd.Series, np.bool_) - - eq = check(assert_type(ts == c_timestamp, bool), bool) - ne = check(assert_type(ts != c_timestamp, bool), bool) - assert eq != ne - - eq = check(assert_type(ts == c_np_dt64, bool), bool) - ne = check(assert_type(ts != c_np_dt64, bool), bool) - assert eq != ne - - eq = check(assert_type(ts == c_dt_datetime, bool), bool) - ne = check(assert_type(ts != c_dt_datetime, bool), bool) - assert eq != ne - - eq_arr = check( - assert_type(ts == c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_arr = check( - assert_type(ts != c_datetimeindex, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_arr != ne_arr).all() - eq_arr = check( - assert_type(ts == c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_arr = check( - assert_type(ts != c_unknown_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_arr != ne_arr).all() - - eq_arr = check( - assert_type(ts == c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_arr = check( - assert_type(ts != c_np_ndarray_dt64, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_arr != ne_arr).all() - - eq_s = check( - assert_type(ts == c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(ts != c_series_timestamp, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - eq_s = check( - assert_type(ts == c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(ts != c_series_dt64, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - -def test_timestamp_eq_ne_rhs() -> None: - # These test equality using the LHS objects __eq__ and __ne__ methods - # The tests are retained for completeness, but are not strictly necessary + ts2 = pd.Timestamp(year=2025, month=8, day=16) + py_dt = dt.datetime(year=2000, month=1, day=1) + np_dt = np.datetime64(1, "ns") + + # >, <= + gt1 = check(assert_type(ts > ts2, bool), bool) + le1 = check(assert_type(ts <= ts2, bool), bool) + assert gt1 != le1 + gt2 = check(assert_type(ts > py_dt, bool), bool) + le2 = check(assert_type(ts <= py_dt, bool), bool) + assert gt2 != le2 + gt3 = check(assert_type(ts > np_dt, bool), bool) + le3 = check(assert_type(ts <= np_dt, bool), bool) + assert gt3 != le3 + gt4 = check(assert_type(py_dt > ts, bool), bool) + le4 = check(assert_type(py_dt <= ts, bool), bool) + assert gt4 != le4 + gt5 = check(assert_type(np_dt > ts, np.bool), bool) + le5 = check(assert_type(np_dt <= ts, np.bool), bool) + assert gt5 != le5 + + # <, >= + lt1 = check(assert_type(ts < ts2, bool), bool) + ge1 = check(assert_type(ts >= ts2, bool), bool) + assert ge1 != lt1 + lt2 = check(assert_type(ts < py_dt, bool), bool) + ge2 = check(assert_type(ts >= py_dt, bool), bool) + assert ge2 != lt2 + lt3 = check(assert_type(ts < np_dt, bool), bool) + ge3 = check(assert_type(ts >= np_dt, bool), bool) + assert ge3 != lt3 + lt4 = check(assert_type(py_dt < ts, bool), bool) + ge4 = check(assert_type(py_dt >= ts, bool), bool) + assert ge4 != lt4 + lt5 = check(assert_type(np_dt < ts, np.bool), bool) + ge5 = check(assert_type(np_dt >= ts, np.bool), bool) + assert ge5 != lt5 + + # =, != + eq1 = check(assert_type(ts == ts2, bool), bool) + ne1 = check(assert_type(ts != ts2, bool), bool) + assert eq1 != ne1 + eq2 = check(assert_type(ts == py_dt, bool), bool) + ne2 = check(assert_type(ts != py_dt, bool), bool) + assert eq2 != ne2 + eq3 = check(assert_type(ts == np_dt, bool), bool) + ne3 = check(assert_type(ts != np_dt, bool), bool) + assert eq3 != ne3 + eq4 = check(assert_type(ts == 1, Literal[False]), bool) + ne4 = check(assert_type(ts != 1, Literal[True]), bool) + assert eq4 != ne4 + eq5 = check(assert_type(ts == (3 + 2j), Literal[False]), bool) + ne5 = check(assert_type(ts != (3 + 2j), Literal[True]), bool) + assert eq5 != ne5 + + # ==, != (ts on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(py_dt == ts, bool), bool) + ne_rhs1 = check(assert_type(py_dt != ts, bool), bool) + assert eq_rhs1 != ne_rhs1 + eq_rhs2 = check(assert_type(np_dt == ts, Any), bool) + ne_rhs2 = check(assert_type(np_dt != ts, Any), bool) + assert eq_rhs2 != ne_rhs2 + + +def test_timestamp_cmp_series() -> None: ts = pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27) - - np_dt64_arr: npt.NDArray[np.datetime64] = np.array( - [1, 2, 3], dtype="datetime64[ns]" - ) - - c_np_dt64 = np.datetime64(1, "ns") - c_dt_datetime = dt.datetime(year=2000, month=1, day=1) - c_datetimeindex = pd.DatetimeIndex(["2000-1-1"]) - # DatetimeIndex, but the type checker detects it to be UnknownIndex. - c_unknown_index = pd.DataFrame({"a": [1]}, index=c_datetimeindex).index - c_np_ndarray_dt64 = np_dt64_arr - c_series_dt64: pd.Series[pd.Timestamp] = pd.Series( - [1, 2, 3], dtype="datetime64[ns]" - ) - - eq_a = check(assert_type(c_np_dt64 == ts, Any), bool) - ne_a = check(assert_type(c_np_dt64 != ts, Any), bool) - assert eq_a != ne_a - - eq = check(assert_type(c_dt_datetime == ts, bool), bool) - ne = check(assert_type(c_dt_datetime != ts, bool), bool) - assert eq != ne - - eq_arr = check( - assert_type(c_datetimeindex == ts, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_arr = check( - assert_type(c_datetimeindex != ts, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_arr != ne_arr).all() - eq_arr = check( - assert_type(c_unknown_index == ts, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_arr = check( - assert_type(c_unknown_index != ts, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_arr != ne_arr).all() - - eq_a = check(assert_type(c_np_ndarray_dt64 != ts, Any), np.ndarray, np.bool_) - ne_a = check(assert_type(c_np_ndarray_dt64 == ts, Any), np.ndarray, np.bool_) - assert (eq_a != ne_a).all() - - eq_s = check( - assert_type(c_series_dt64 == ts, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(c_series_dt64 != ts, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() + ts_ser = pd.Series(pd.DatetimeIndex(["2000-1-1", "2000-1-2"])) + check(assert_type(ts_ser, TimestampSeries), pd.Series, pd.Timestamp) + + # >, <= + gt1 = check(assert_type(ts > ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + le1 = check(assert_type(ts <= ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt1 != le1).all() + gt2 = check(assert_type(ts_ser > ts, "pd.Series[bool]"), pd.Series, np.bool) + le2 = check(assert_type(ts_ser <= ts, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt2 != le2).all() + + # <, >= + lt1 = check(assert_type(ts < ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + ge1 = check(assert_type(ts >= ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt1 != ge1).all() + lt2 = check(assert_type(ts_ser < ts, "pd.Series[bool]"), pd.Series, np.bool) + ge2 = check(assert_type(ts_ser >= ts, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt2 != ge2).all() + + # ==, != + eq1 = check(assert_type(ts == ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + ne1 = check(assert_type(ts != ts_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq1 != ne1).all() + + # ==, != (ts on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(ts_ser == ts, "pd.Series[bool]"), pd.Series, np.bool) + ne_rhs1 = check(assert_type(ts_ser != ts, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq_rhs1 != ne_rhs1).all() + + +def test_timestamp_cmp_index() -> None: + ts = pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27) + dt_idx = pd.DatetimeIndex(["2000-1-1"]) + # DatetimeIndex, but the type checker thinks it is UnknownIndex. + un_idx = pd.DataFrame({"a": [1]}, index=dt_idx).index + + # >, <= + gt_dt1 = check(assert_type(ts > dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + le_dt1 = check(assert_type(ts <= dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_dt1 != le_dt1).all() + gt_dt2 = check(assert_type(dt_idx > ts, np_1darray[np.bool]), np_1darray[np.bool]) + le_dt2 = check(assert_type(dt_idx <= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_dt2 != le_dt2).all() + gt_un1 = check(assert_type(ts > un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + le_un1 = check(assert_type(ts <= un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_un1 != le_un1).all() + gt_un2 = check(assert_type(un_idx > ts, np_1darray[np.bool]), np_1darray[np.bool]) + le_un2 = check(assert_type(un_idx <= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_un2 != le_un2).all() + + # <, >= + lt_dt1 = check(assert_type(ts < dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ge_dt1 = check(assert_type(ts >= dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_dt1 != ge_dt1).all() + lt_dt2 = check(assert_type(dt_idx < ts, np_1darray[np.bool]), np_1darray[np.bool]) + ge_dt2 = check(assert_type(dt_idx >= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_dt2 != ge_dt2).all() + lt_un1 = check(assert_type(ts < un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ge_un1 = check(assert_type(ts >= un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_un1 != ge_un1).all() + lt_un2 = check(assert_type(un_idx < ts, np_1darray[np.bool]), np_1darray[np.bool]) + ge_un2 = check(assert_type(un_idx >= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_un2 != ge_un2).all() + + # ==, != + eq_dt1 = check(assert_type(ts == dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ne_dt1 = check(assert_type(ts != dt_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_dt1 != ne_dt1).all() + eq_un1 = check(assert_type(ts == un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ne_un1 = check(assert_type(ts != un_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_un1 != ne_un1).all() + + # ==, != (ts on the rhs, use == and != of lhs) + eq_rhs_dt1 = check( + assert_type(dt_idx == ts, np_1darray[np.bool]), np_1darray[np.bool] + ) + ne_rhs_dt1 = check( + assert_type(dt_idx != ts, np_1darray[np.bool]), np_1darray[np.bool] + ) + assert (eq_rhs_dt1 != ne_rhs_dt1).all() + eq_rhs_un1 = check( + assert_type(un_idx == ts, np_1darray[np.bool]), np_1darray[np.bool] + ) + ne_rhs_un1 = check( + assert_type(un_idx != ts, np_1darray[np.bool]), np_1darray[np.bool] + ) + assert (eq_rhs_un1 != ne_rhs_un1).all() + + +def test_timestamp_cmp_array() -> None: + ts = pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27) + arr_1d = pd.to_datetime([1, 2, 3, 4]).to_numpy() + arr_2d = arr_1d.reshape(2, 2) + arr_nd = arr_1d.reshape([4]) + + # >, <= + gt_nd1 = check(assert_type(ts > arr_nd, np_ndarray_bool), np.ndarray, np.bool) + le_nd1 = check(assert_type(ts <= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd1 != le_nd1).all() + gt_nd2 = check(assert_type(arr_nd > ts, np_ndarray_bool), np.ndarray, np.bool) + le_nd2 = check(assert_type(arr_nd <= ts, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd2 != le_nd2).all() + gt_2d1 = check(assert_type(ts > arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + le_2d1 = check(assert_type(ts <= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (gt_2d1 != le_2d1).all() + gt_2d2 = check(assert_type(arr_2d > ts, np_2darray[np.bool]), np_2darray[np.bool]) + le_2d2 = check(assert_type(arr_2d <= ts, np_2darray[np.bool]), np_2darray[np.bool]) + assert (gt_2d2 != le_2d2).all() + gt_1d1 = check(assert_type(ts > arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + le_1d1 = check(assert_type(ts <= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_1d1 != le_1d1).all() + gt_1d2 = check(assert_type(arr_1d > ts, np_1darray[np.bool]), np_1darray[np.bool]) + le_1d2 = check(assert_type(arr_1d <= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_1d2 != le_1d2).all() + + # <, >= + lt_nd1 = check(assert_type(ts < arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ge_nd1 = check(assert_type(ts >= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd1 != ge_nd1).all() + lt_nd2 = check(assert_type(arr_nd < ts, np_ndarray_bool), np.ndarray, np.bool) + ge_nd2 = check(assert_type(arr_nd >= ts, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd2 != ge_nd2).all() + lt_2d1 = check(assert_type(ts < arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ge_2d1 = check(assert_type(ts >= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (lt_2d1 != ge_2d1).all() + lt_2d2 = check(assert_type(arr_2d < ts, np_2darray[np.bool]), np_2darray[np.bool]) + ge_2d2 = check(assert_type(arr_2d >= ts, np_2darray[np.bool]), np_2darray[np.bool]) + assert (lt_2d2 != ge_2d2).all() + lt_1d1 = check(assert_type(ts < arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ge_1d1 = check(assert_type(ts >= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_1d1 != ge_1d1).all() + lt_1d2 = check(assert_type(arr_1d < ts, np_1darray[np.bool]), np_1darray[np.bool]) + ge_1d2 = check(assert_type(arr_1d >= ts, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_1d2 != ge_1d2).all() + + # ==, != + eq_nd1 = check(assert_type(ts == arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ne_nd1 = check(assert_type(ts != arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (eq_nd1 != ne_nd1).all() + eq_2d1 = check(assert_type(ts == arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ne_2d1 = check(assert_type(ts != arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (eq_2d1 != ne_2d1).all() + eq_1d1 = check(assert_type(ts == arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ne_1d1 = check(assert_type(ts != arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_1d1 != ne_1d1).all() + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs_nd1 = check(assert_type(arr_nd == ts, Any), np_ndarray_bool) + ne_rhs_nd1 = check(assert_type(arr_nd != ts, Any), np_ndarray_bool) + assert (eq_rhs_nd1 != ne_rhs_nd1).all() + eq_rhs_2d1 = check(assert_type(arr_2d == ts, Any), np_2darray[np.bool]) + ne_rhs_2d1 = check(assert_type(arr_2d != ts, Any), np_2darray[np.bool]) + assert (eq_rhs_2d1 != ne_rhs_2d1).all() + eq_rhs_1d1 = check(assert_type(arr_1d == ts, Any), np_1darray[np.bool]) + ne_rhs_1d1 = check(assert_type(arr_1d != ts, Any), np_1darray[np.bool]) + assert (eq_rhs_1d1 != ne_rhs_1d1).all() def test_timestamp_types_init() -> None: @@ -1914,124 +1862,160 @@ def test_period_add_subtract() -> None: check(assert_type(as_period_index - p, pd.Index), pd.Index) -def test_period_cmp() -> None: +def test_period_cmp_scalar() -> None: p = pd.Period("2012-1-1", freq="D") + p2 = pd.Period("2012-1-2", freq="D") - c_period = pd.Period("2012-1-1", freq="D") - c_period_index = pd.period_range("2012-1-1", periods=10, freq="D") - c_period_series = pd.Series(c_period_index) + # >, <= + gt1 = check(assert_type(p > p2, bool), bool) + le1 = check(assert_type(p <= p2, bool), bool) + assert gt1 != le1 - eq = check(assert_type(p == c_period, bool), bool) - ne = check(assert_type(p != c_period, bool), bool) - assert eq != ne + # <, >= + lt1 = check(assert_type(p < p2, bool), bool) + ge1 = check(assert_type(p >= p2, bool), bool) + assert lt1 != ge1 - eq_a = check( - assert_type(p == c_period_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_q = check( - assert_type(p != c_period_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_a != ne_q).all() - - eq_s = check( - assert_type(p == c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(p != c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() - - eq = check(assert_type(c_period == p, bool), bool) - ne = check(assert_type(c_period != p, bool), bool) - assert eq != ne - - eq_a = check( - assert_type(c_period_index == p, np_ndarray_bool), np.ndarray, np.bool_ - ) - ne_a = check( - assert_type(c_period_index != p, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (eq_a != ne_a).all() - - eq_s = check( - assert_type(c_period_series == p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ne_s = check( - assert_type(c_period_series != p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (eq_s != ne_s).all() + # ==, != + eq1 = check(assert_type(p == p2, bool), bool) + ne1 = check(assert_type(p != p2, bool), bool) + assert eq1 != ne1 + eq2 = check(assert_type(p == 1, Literal[False]), bool) + ne2 = check(assert_type(p != 1, Literal[True]), bool) + assert eq2 != ne2 - gt = check(assert_type(p > c_period, bool), bool) - le = check(assert_type(p <= c_period, bool), bool) - assert gt != le - gt_a = check(assert_type(p > c_period_index, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check( - assert_type(p <= c_period_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - gt_s = check( - assert_type(p > c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - le_s = check( - assert_type(p <= c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (gt_s != le_s).all() - - gt = check(assert_type(c_period > p, bool), bool) - le = check(assert_type(c_period <= p, bool), bool) - assert gt != le - - gt_a = check(assert_type(c_period_index > p, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check( - assert_type(c_period_index <= p, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (gt_a != le_a).all() - - gt_s = check( - assert_type(c_period_series > p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - le_s = check( - assert_type(c_period_series <= p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (gt_s != le_s).all() - - lt = check(assert_type(p < c_period, bool), bool) - ge = check(assert_type(p >= c_period, bool), bool) - assert lt != ge - - lt_a = check(assert_type(p < c_period_index, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check( - assert_type(p >= c_period_index, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - lt_s = check( - assert_type(p < c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ge_s = check( - assert_type(p >= c_period_series, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (lt_s != ge_s).all() - - lt = check(assert_type(c_period < p, bool), bool) - ge = check(assert_type(c_period >= p, bool), bool) - assert lt != ge - - lt_a = check(assert_type(c_period_index < p, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check( - assert_type(c_period_index >= p, np_ndarray_bool), np.ndarray, np.bool_ - ) - assert (lt_a != ge_a).all() - - lt_s = check( - assert_type(c_period_series < p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - ge_s = check( - assert_type(c_period_series >= p, "pd.Series[bool]"), pd.Series, np.bool_ - ) - assert (lt_s != ge_s).all() +def test_period_cmp_series() -> None: + p = pd.Period("2012-1-1", freq="D") + p_ser = pd.Series(pd.period_range("2012-1-1", periods=10, freq="D")) + + # >, <= + gt1 = check(assert_type(p > p_ser, "pd.Series[bool]"), pd.Series, np.bool) + le1 = check(assert_type(p <= p_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt1 != le1).all() + gt2 = check(assert_type(p_ser > p, "pd.Series[bool]"), pd.Series, np.bool) + le2 = check(assert_type(p_ser <= p, "pd.Series[bool]"), pd.Series, np.bool) + assert (gt2 != le2).all() + + # <, >= + lt1 = check(assert_type(p < p_ser, "pd.Series[bool]"), pd.Series, np.bool) + ge1 = check(assert_type(p >= p_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt1 != ge1).all() + lt2 = check(assert_type(p_ser < p, "pd.Series[bool]"), pd.Series, np.bool) + ge2 = check(assert_type(p_ser >= p, "pd.Series[bool]"), pd.Series, np.bool) + assert (lt2 != ge2).all() + + # ==, != + eq1 = check(assert_type(p == p_ser, "pd.Series[bool]"), pd.Series, np.bool) + ne1 = check(assert_type(p != p_ser, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq1 != ne1).all() + + # ==, != (p on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(p_ser == p, "pd.Series[bool]"), pd.Series, np.bool) + ne_rhs1 = check(assert_type(p_ser != p, "pd.Series[bool]"), pd.Series, np.bool) + assert (eq_rhs1 != ne_rhs1).all() + + +def test_period_cmp_index() -> None: + p = pd.Period("2012-1-1", freq="D") + p_idx = pd.period_range("2012-1-1", periods=10, freq="D") + + # >, <= + gt1 = check(assert_type(p > p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + le1 = check(assert_type(p <= p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt1 != le1).all() + gt2 = check(assert_type(p_idx > p, np_1darray[np.bool]), np_1darray[np.bool]) + le2 = check(assert_type(p_idx <= p, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt2 != le2).all() + + # <, >= + lt1 = check(assert_type(p < p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ge1 = check(assert_type(p >= p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt1 != ge1).all() + lt2 = check(assert_type(p_idx < p, np_1darray[np.bool]), np_1darray[np.bool]) + ge2 = check(assert_type(p_idx >= p, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt2 != ge2).all() + + # ==, != + eq1 = check(assert_type(p == p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + ne1 = check(assert_type(p != p_idx, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq1 != ne1).all() + + # ==, != (p on the rhs, use == and != of lhs) + eq_rhs1 = check(assert_type(p_idx == p, np_1darray[np.bool]), np_1darray[np.bool]) + ne_rhs1 = check(assert_type(p_idx != p, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_rhs1 != ne_rhs1).all() + + +def test_period_cmp_array() -> None: + p = pd.Period("2012-1-1", freq="D") + arr_1d = pd.period_range("2012-1-1", periods=4, freq="D").to_numpy() + arr_2d = arr_1d.reshape(2, 2) + arr_nd = arr_1d.reshape([4]) + + # >, <= + gt_nd1 = check(assert_type(p > arr_nd, np_ndarray_bool), np.ndarray, np.bool) + le_nd1 = check(assert_type(p <= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd1 != le_nd1).all() + gt_2d1 = check(assert_type(p > arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + le_2d1 = check(assert_type(p <= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (gt_2d1 != le_2d1).all() + gt_1d1 = check(assert_type(p > arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + le_1d1 = check(assert_type(p <= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (gt_1d1 != le_1d1).all() + # p on the rhs, type depends on np.ndarray > and <= methods + gt_nd2 = check(assert_type(arr_nd > p, np_ndarray_bool), np.ndarray, np.bool) + le_nd2 = check(assert_type(arr_nd <= p, np_ndarray_bool), np.ndarray, np.bool) + assert (gt_nd2 != le_nd2).all() + gt_2d2 = check(assert_type(arr_2d > p, np_ndarray_bool), np_2darray[np.bool]) + le_2d2 = check(assert_type(arr_2d <= p, np_ndarray_bool), np_2darray[np.bool]) + assert (gt_2d2 != le_2d2).all() + gt_1d2 = check(assert_type(arr_1d > p, np_ndarray_bool), np_1darray[np.bool]) + le_1d2 = check(assert_type(arr_1d <= p, np_ndarray_bool), np_1darray[np.bool]) + assert (gt_1d2 != le_1d2).all() + + # <, >= + lt_nd1 = check(assert_type(p < arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ge_nd1 = check(assert_type(p >= arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd1 != ge_nd1).all() + lt_2d1 = check(assert_type(p < arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ge_2d1 = check(assert_type(p >= arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (lt_2d1 != ge_2d1).all() + lt_1d1 = check(assert_type(p < arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ge_1d1 = check(assert_type(p >= arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (lt_1d1 != ge_1d1).all() + # p on the rhs, type depends on np.ndarray < and >= methods + lt_nd2 = check(assert_type(arr_nd < p, np_ndarray_bool), np.ndarray, np.bool) + ge_nd2 = check(assert_type(arr_nd >= p, np_ndarray_bool), np.ndarray, np.bool) + assert (lt_nd2 != ge_nd2).all() + lt_2d2 = check(assert_type(arr_2d < p, np_ndarray_bool), np_2darray[np.bool]) + ge_2d2 = check(assert_type(arr_2d >= p, np_ndarray_bool), np_2darray[np.bool]) + assert (lt_2d2 != ge_2d2).all() + lt_1d2 = check(assert_type(arr_1d < p, np_ndarray_bool), np_1darray[np.bool]) + ge_1d2 = check(assert_type(arr_1d >= p, np_ndarray_bool), np_1darray[np.bool]) + assert (lt_1d2 != ge_1d2).all() + + # ==, != + eq_nd1 = check(assert_type(p == arr_nd, np_ndarray_bool), np.ndarray, np.bool) + ne_nd1 = check(assert_type(p != arr_nd, np_ndarray_bool), np.ndarray, np.bool) + assert (eq_nd1 != ne_nd1).all() + eq_2d1 = check(assert_type(p == arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + ne_2d1 = check(assert_type(p != arr_2d, np_2darray[np.bool]), np_2darray[np.bool]) + assert (eq_2d1 != ne_2d1).all() + eq_1d1 = check(assert_type(p == arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + ne_1d1 = check(assert_type(p != arr_1d, np_1darray[np.bool]), np_1darray[np.bool]) + assert (eq_1d1 != ne_1d1).all() + + # ==, != (td on the rhs, use == and != of lhs) + eq_rhs_nd1 = check(assert_type(arr_nd == p, Any), np_ndarray_bool) + ne_rhs_nd1 = check(assert_type(arr_nd != p, Any), np_ndarray_bool) + assert (eq_rhs_nd1 != ne_rhs_nd1).all() + eq_rhs_2d1 = check(assert_type(arr_2d == p, Any), np_2darray[np.bool]) + ne_rhs_2d1 = check(assert_type(arr_2d != p, Any), np_2darray[np.bool]) + assert (eq_rhs_2d1 != ne_rhs_2d1).all() + eq_rhs_1d1 = check(assert_type(arr_1d == p, Any), np_1darray[np.bool]) + ne_rhs_1d1 = check(assert_type(arr_1d != p, Any), np_1darray[np.bool]) + assert (eq_rhs_1d1 != ne_rhs_1d1).all() def test_period_methods() -> None: diff --git a/tests/test_string_accessors.py b/tests/test_string_accessors.py index 5796ee40d..b501f681c 100644 --- a/tests/test_string_accessors.py +++ b/tests/test_string_accessors.py @@ -10,7 +10,7 @@ PD_LTE_23, TYPE_CHECKING_INVALID_USAGE, check, - np_ndarray_bool, + np_1darray, ) DATA = ["applep", "bananap", "Cherryp", "DATEp", "eGGpLANTp", "123p", "23.45p"] @@ -71,40 +71,42 @@ def test_string_accessors_boolean_series(): def test_string_accessors_boolean_index(): idx = pd.Index(DATA) - _check = functools.partial(check, klass=np.ndarray, dtype=np.bool_) - _check(assert_type(idx.str.startswith("a"), np_ndarray_bool)) + _check = functools.partial(check, klass=np_1darray[np.bool]) + _check(assert_type(idx.str.startswith("a"), np_1darray[np.bool])) _check( - assert_type(idx.str.startswith(("a", "b")), np_ndarray_bool), + assert_type(idx.str.startswith(("a", "b")), np_1darray[np.bool]), ) _check( - assert_type(idx.str.contains("a"), np_ndarray_bool), + assert_type(idx.str.contains("a"), np_1darray[np.bool]), ) if PD_LTE_23: # Bug in pandas 3.0 dev https://github.com/pandas-dev/pandas/issues/61942 _check( assert_type( - idx.str.contains(re.compile(r"a"), regex=True), np_ndarray_bool + idx.str.contains(re.compile(r"a"), regex=True), np_1darray[np.bool] ), ) - _check(assert_type(idx.str.endswith("e"), np_ndarray_bool)) - _check(assert_type(idx.str.endswith(("e", "f")), np_ndarray_bool)) - _check(assert_type(idx.str.fullmatch("apple"), np_ndarray_bool)) + _check(assert_type(idx.str.endswith("e"), np_1darray[np.bool])) + _check(assert_type(idx.str.endswith(("e", "f")), np_1darray[np.bool])) + _check(assert_type(idx.str.fullmatch("apple"), np_1darray[np.bool])) if PD_LTE_23: # Bug in 3.0 dev: https://github.com/pandas-dev/pandas/issues/61952 - _check(assert_type(idx.str.fullmatch(re.compile(r"apple")), np_ndarray_bool)) - _check(assert_type(idx.str.isalnum(), np_ndarray_bool)) - _check(assert_type(idx.str.isalpha(), np_ndarray_bool)) - _check(assert_type(idx.str.isdecimal(), np_ndarray_bool)) - _check(assert_type(idx.str.isdigit(), np_ndarray_bool)) - _check(assert_type(idx.str.isnumeric(), np_ndarray_bool)) - _check(assert_type(idx.str.islower(), np_ndarray_bool)) - _check(assert_type(idx.str.isspace(), np_ndarray_bool)) - _check(assert_type(idx.str.istitle(), np_ndarray_bool)) - _check(assert_type(idx.str.isupper(), np_ndarray_bool)) - _check(assert_type(idx.str.match("pp"), np_ndarray_bool)) + _check( + assert_type(idx.str.fullmatch(re.compile(r"apple")), np_1darray[np.bool]) + ) + _check(assert_type(idx.str.isalnum(), np_1darray[np.bool])) + _check(assert_type(idx.str.isalpha(), np_1darray[np.bool])) + _check(assert_type(idx.str.isdecimal(), np_1darray[np.bool])) + _check(assert_type(idx.str.isdigit(), np_1darray[np.bool])) + _check(assert_type(idx.str.isnumeric(), np_1darray[np.bool])) + _check(assert_type(idx.str.islower(), np_1darray[np.bool])) + _check(assert_type(idx.str.isspace(), np_1darray[np.bool])) + _check(assert_type(idx.str.istitle(), np_1darray[np.bool])) + _check(assert_type(idx.str.isupper(), np_1darray[np.bool])) + _check(assert_type(idx.str.match("pp"), np_1darray[np.bool])) if PD_LTE_23: # Bug in 3.0 dev: https://github.com/pandas-dev/pandas/issues/61952 - _check(assert_type(idx.str.match(re.compile(r"pp")), np_ndarray_bool)) + _check(assert_type(idx.str.match(re.compile(r"pp")), np_1darray[np.bool])) def test_string_accessors_integer_series(): diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 7e39e4a75..356d1516a 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -18,7 +18,6 @@ WE, ) import numpy as np -from numpy import typing as npt import pandas as pd from pandas.api.typing import NaTType from pandas.core.tools.datetimes import FulldatetimeDict @@ -34,6 +33,7 @@ PD_LTE_23, TYPE_CHECKING_INVALID_USAGE, check, + np_1darray, pytest_warns_bounded, ) @@ -50,9 +50,13 @@ ) if TYPE_CHECKING: + from pandas.core.series import ( + IntervalSeries, + OffsetSeries, + TimestampSeries, + ) from pandas.core.series import PeriodSeries # noqa: F401 from pandas.core.series import TimedeltaSeries # noqa: F401 - from pandas.core.series import TimestampSeries else: TimestampSeries: TypeAlias = pd.Series @@ -61,8 +65,6 @@ else: Pandas4Warning: TypeAlias = FutureWarning # type: ignore[no-redef] -from tests import np_ndarray_bool - def test_types_init() -> None: check(assert_type(pd.Timestamp("2021-03-01T12"), pd.Timestamp), pd.Timestamp) @@ -350,12 +352,12 @@ def test_comparisons_datetimeindex() -> None: # GH 74 dti = pd.date_range("2000-01-01", "2000-01-10") ts = pd.Timestamp("2000-01-05") - check(assert_type((dti < ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti > ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti >= ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti <= ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti == ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti != ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti < ts), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type((dti > ts), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type((dti >= ts), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type((dti <= ts), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type((dti == ts), np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type((dti != ts), np_1darray[np.bool]), np_1darray[np.bool]) def test_to_datetime_nat() -> None: @@ -427,8 +429,8 @@ def test_series_dt_accessors() -> None: upper="2.3.99", ): check( - assert_type(s0.dt.to_pydatetime(), np.ndarray), - np.ndarray if PD_LTE_23 else pd.Series, + assert_type(s0.dt.to_pydatetime(), np_1darray[np.object_]), + np_1darray[np.object_] if PD_LTE_23 else pd.Series, dt.datetime, ) s0_local = s0.dt.tz_localize("UTC") @@ -560,7 +562,11 @@ def test_series_dt_accessors() -> None: upper="3.0.99", ), ): - check(assert_type(s2.dt.to_pytimedelta(), np.ndarray), np.ndarray) + check( + assert_type(s2.dt.to_pytimedelta(), np_1darray[np.object_]), + np_1darray[np.object_], + dt.timedelta, + ) check(assert_type(s2.dt.total_seconds(), "pd.Series[float]"), pd.Series, float) check(assert_type(s2.dt.unit, TimeUnit), str) check(assert_type(s2.dt.as_unit("s"), "TimedeltaSeries"), pd.Series, pd.Timedelta) @@ -601,9 +607,11 @@ def test_datetimeindex_accessors() -> None: i0 = pd.date_range(start="2022-06-01", periods=10) check(assert_type(i0, pd.DatetimeIndex), pd.DatetimeIndex, pd.Timestamp) - check(assert_type(i0.date, np.ndarray), np.ndarray, dt.date) - check(assert_type(i0.time, np.ndarray), np.ndarray, dt.time) - check(assert_type(i0.timetz, np.ndarray), np.ndarray, dt.time) + check(assert_type(i0.date, np_1darray[np.object_]), np_1darray[np.object_], dt.date) + check(assert_type(i0.time, np_1darray[np.object_]), np_1darray[np.object_], dt.time) + check( + assert_type(i0.timetz, np_1darray[np.object_]), np_1darray[np.object_], dt.time + ) check(assert_type(i0.year, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.month, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.day, "pd.Index[int]"), pd.Index, np.int32) @@ -618,13 +626,13 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.dayofyear, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.day_of_year, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.quarter, "pd.Index[int]"), pd.Index, np.int32) - check(assert_type(i0.is_month_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_month_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_quarter_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_quarter_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_year_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_year_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.is_leap_year, npt.NDArray[np.bool_]), np.ndarray, np.bool_) + check(assert_type(i0.is_month_start, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_month_end, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_quarter_start, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_quarter_end, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_year_start, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_year_end, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(i0.is_leap_year, np_1darray[np.bool]), np_1darray[np.bool]) check(assert_type(i0.daysinmonth, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.days_in_month, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.tz, Optional[dt.tzinfo]), type(None)) @@ -632,8 +640,8 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.isocalendar(), pd.DataFrame), pd.DataFrame) check(assert_type(i0.to_period("D"), pd.PeriodIndex), pd.PeriodIndex, pd.Period) check( - assert_type(i0.to_pydatetime(), npt.NDArray[np.object_]), - np.ndarray, + assert_type(i0.to_pydatetime(), np_1darray[np.object_]), + np_1darray[np.object_], dt.datetime, ) ilocal = i0.tz_localize("UTC") @@ -674,7 +682,11 @@ def test_timedeltaindex_accessors() -> None: check(assert_type(i0.microseconds, pd.Index), pd.Index, np.integer) check(assert_type(i0.nanoseconds, pd.Index), pd.Index, np.integer) check(assert_type(i0.components, pd.DataFrame), pd.DataFrame) - check(assert_type(i0.to_pytimedelta(), np.ndarray), np.ndarray) + check( + assert_type(i0.to_pytimedelta(), np_1darray[np.object_]), + np_1darray[np.object_], + dt.timedelta, + ) check(assert_type(i0.total_seconds(), pd.Index), pd.Index, float) check( assert_type(i0.round("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta @@ -875,11 +887,223 @@ def test_timestampseries_offset() -> None: check(assert_type(shifted_vv, pd.DatetimeIndex), pd.DatetimeIndex) -def test_types_to_numpy() -> None: +def test_series_types_to_numpy() -> None: td_s = pd.to_timedelta(pd.Series([10, 20]), "minutes") - check(assert_type(td_s.to_numpy(), np.ndarray), np.ndarray) - check(assert_type(td_s.to_numpy(dtype="int", copy=True), np.ndarray), np.ndarray) - check(assert_type(td_s.to_numpy(na_value=pd.Timedelta(0)), np.ndarray), np.ndarray) + ts_s = pd.to_datetime(pd.Series(["2020-01-01", "2020-01-02"])) + p_s = pd.Series(pd.period_range("2012-1-1", periods=10, freq="D")) + o_s = cast( + "OffsetSeries", pd.Series([pd.DateOffset(days=1), pd.DateOffset(days=2)]) + ) + i_s = cast("IntervalSeries", pd.interval_range(1, 2).to_series()) + + # default dtype + check( + assert_type(td_s.to_numpy(), np_1darray[np.timedelta64]), + np_1darray, + dtype=np.timedelta64, + ) + check( + assert_type( + td_s.to_numpy(na_value=pd.Timedelta(0)), np_1darray[np.timedelta64] + ), + np_1darray, + dtype=np.timedelta64, + ) + check( + assert_type(ts_s.to_numpy(), np_1darray[np.datetime64]), + np_1darray, + dtype=np.datetime64, + ) + check( + assert_type(ts_s.to_numpy(na_value=pd.Timestamp(1)), np_1darray[np.datetime64]), + np_1darray, + dtype=np.datetime64, + ) + check( + assert_type(p_s.to_numpy(), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Period, + ) + check( + assert_type(p_s.to_numpy(na_value=pd.Timestamp(1)), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Period, + ) + check( + assert_type(o_s.to_numpy(), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.DateOffset, + ) + check( + assert_type( + o_s.to_numpy(na_value=pd.Timedelta(days=1)), np_1darray[np.object_] + ), + np_1darray[np.object_], + dtype=pd.DateOffset, + ) + check( + assert_type(i_s.to_numpy(), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Interval, + ) + check( + assert_type( + i_s.to_numpy(na_value=pd.Timedelta(days=1)), np_1darray[np.object_] + ), + np_1darray[np.object_], + dtype=pd.Interval, + ) + + # passed dtype-like with statically unknown generic + check( + assert_type(td_s.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(ts_s.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(p_s.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(o_s.to_numpy(dtype="bytes", copy=True), np_1darray), + np_1darray, + dtype=np.bytes_, + ) + check( + assert_type(i_s.to_numpy(dtype="bytes", copy=True), np_1darray), + np_1darray, + dtype=np.bytes_, + ) + + # passed dtype-like with statically known generic + check( + assert_type(td_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(ts_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(p_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(o_s.to_numpy(dtype=np.bytes_), np_1darray[np.bytes_]), + np_1darray, + dtype=np.bytes_, + ) + check( + assert_type(i_s.to_numpy(dtype=np.bytes_), np_1darray[np.bytes_]), + np_1darray, + dtype=np.bytes_, + ) + + +def test_index_types_to_numpy() -> None: + td_i = pd.timedelta_range(10, 20) + ts_i = pd.date_range("2025-1-1", "2025-1-2") + p_i = pd.period_range("2025-1-1", periods=10, freq="D") + i_i = pd.interval_range(1, 2) + + # default dtype + check( + assert_type(td_i.to_numpy(), np_1darray[np.timedelta64]), + np_1darray, + dtype=np.timedelta64, + ) + check( + assert_type( + td_i.to_numpy(na_value=pd.Timedelta(0)), np_1darray[np.timedelta64] + ), + np_1darray, + dtype=np.timedelta64, + ) + check( + assert_type(ts_i.to_numpy(), np_1darray[np.datetime64]), + np_1darray, + dtype=np.datetime64, + ) + check( + assert_type(ts_i.to_numpy(na_value=pd.Timestamp(1)), np_1darray[np.datetime64]), + np_1darray, + dtype=np.datetime64, + ) + check( + assert_type(p_i.to_numpy(), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Period, + ) + check( + assert_type(p_i.to_numpy(na_value=pd.Timestamp(1)), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Period, + ) + check( + assert_type(i_i.to_numpy(), np_1darray[np.object_]), + np_1darray[np.object_], + dtype=pd.Interval, + ) + check( + assert_type( + i_i.to_numpy(na_value=pd.Timedelta(days=1)), np_1darray[np.object_] + ), + np_1darray[np.object_], + dtype=pd.Interval, + ) + + # passed dtype-like with statically unknown generic + check( + assert_type(td_i.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(ts_i.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(p_i.to_numpy(dtype="int", copy=True), np_1darray), + np_1darray, + dtype=np.integer, + ) + check( + assert_type(i_i.to_numpy(dtype="bytes", copy=True), np_1darray), + np_1darray, + dtype=np.bytes_, + ) + + # passed dtype-like with statically known generic + check( + assert_type(td_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(ts_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(p_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + np_1darray, + dtype=np.int64, + ) + check( + assert_type(i_i.to_numpy(dtype=np.bytes_), np_1darray[np.bytes_]), + np_1darray, + dtype=np.bytes_, + ) def test_to_timedelta_units() -> None: