diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index 5152f32b9..5913da7da 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -4,9 +4,9 @@ from typing import ( Literal, TypeVar, overload, + type_check_only, ) -import numpy as np from pandas import ( IntervalIndex, Series, @@ -17,8 +17,8 @@ from pandas import ( from pandas._typing import ( IntervalClosedType, IntervalT, - np_1darray, - npt, + np_1darray_bool, + np_ndarray, ) VALID_CLOSED: frozenset[str] @@ -27,6 +27,7 @@ _OrderableScalarT = TypeVar("_OrderableScalarT", bound=int | float) _OrderableTimesT = TypeVar("_OrderableTimesT", bound=Timestamp | Timedelta) _OrderableT = TypeVar("_OrderableT", bound=int | float | Timestamp | Timedelta) +@type_check_only class _LengthDescriptor: @overload def __get__( @@ -37,8 +38,9 @@ class _LengthDescriptor: self, instance: Interval[_OrderableTimesT], owner: Any ) -> Timedelta: ... @overload - def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... + def __get__(self, instance: IntervalMixin, owner: Any) -> np_ndarray: ... +@type_check_only class _MidDescriptor: @overload def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ... @@ -47,7 +49,7 @@ class _MidDescriptor: self, instance: Interval[_OrderableTimesT], owner: Any ) -> _OrderableTimesT: ... @overload - def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... + def __get__(self, instance: IntervalMixin, owner: Any) -> np_ndarray: ... class IntervalMixin: @property @@ -68,8 +70,8 @@ class Interval(IntervalMixin, Generic[_OrderableT]): def right(self: Interval[_OrderableT]) -> _OrderableT: ... @property def closed(self) -> IntervalClosedType: ... - mid: _MidDescriptor - length: _LengthDescriptor + mid = _MidDescriptor() + length = _LengthDescriptor() def __init__( self, left: _OrderableT, @@ -166,9 +168,7 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __gt__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __gt__( - self: IntervalT, other: IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_bool: ... @overload def __gt__( self, @@ -177,9 +177,7 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __lt__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __lt__( - self: IntervalT, other: IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_bool: ... @overload def __lt__( self, @@ -188,9 +186,7 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __ge__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __ge__( - self: IntervalT, other: IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_bool: ... @overload def __ge__( self, @@ -199,15 +195,11 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __le__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __le__( - self: IntervalT, other: IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_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_1darray[np.bool]: ... + def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_bool: ... @overload def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload @@ -215,29 +207,8 @@ 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_1darray[np.bool]: ... + def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_1darray_bool: ... @overload def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload def __ne__(self, other: object) -> Literal[True]: ... - -class IntervalTree(IntervalMixin): - def __init__( - self, - left: np.ndarray, - right: np.ndarray, - closed: IntervalClosedType = ..., - leaf_size: int = ..., - ) -> None: ... - def get_indexer(self, target) -> npt.NDArray[np.intp]: ... - def get_indexer_non_unique( - self, target - ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ... - _na_count: int - @property - def is_overlapping(self) -> bool: ... - @property - def is_monotonic_increasing(self) -> bool: ... - def clear_mapping(self) -> None: ... diff --git a/pandas-stubs/_libs/tslibs/offsets.pyi b/pandas-stubs/_libs/tslibs/offsets.pyi index a6cae3ce3..299feade0 100644 --- a/pandas-stubs/_libs/tslibs/offsets.pyi +++ b/pandas-stubs/_libs/tslibs/offsets.pyi @@ -14,10 +14,14 @@ from typing import ( from dateutil.relativedelta import weekday as WeekdayClass import numpy as np +from numpy import typing as npt from pandas import Timestamp from typing_extensions import Self -from pandas._typing import npt +from pandas._typing import ( + ShapeT, + np_ndarray, +) from pandas.tseries.holiday import AbstractHolidayCalendar @@ -37,7 +41,9 @@ class BaseOffset: @property def base(self) -> BaseOffset: ... @overload - def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... + def __add__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.object_]: ... @overload def __add__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload @@ -47,7 +53,9 @@ class BaseOffset: @overload def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ... @overload - def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ... + def __radd__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.object_]: ... @overload def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload @@ -68,11 +76,15 @@ class BaseOffset: @overload def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ... @overload - def __mul__(self, other: np.ndarray) -> np.ndarray: ... + def __mul__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.object_]: ... @overload def __mul__(self, other: int) -> Self: ... @overload - def __rmul__(self, other: np.ndarray) -> np.ndarray: ... + def __rmul__( + self, other: np_ndarray[ShapeT, np.object_] + ) -> np_ndarray[ShapeT, np.object_]: ... @overload def __rmul__(self, other: int) -> Self: ... def __neg__(self) -> Self: ... diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index f121fc395..f697641be 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -20,7 +20,7 @@ from pandas._libs.tslibs.timestamps import Timestamp from pandas._typing import ( PeriodFrequency, ShapeT, - np_1darray, + np_1darray_bool, np_ndarray, ) @@ -107,7 +107,7 @@ class Period(PeriodMixin): @overload def __eq__(self, other: Self) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __eq__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __eq__(self, other: Index) -> np_1darray_bool: ... # type: ignore[overload-overlap] @overload def __eq__(self, other: Series[Period]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload @@ -117,7 +117,7 @@ class Period(PeriodMixin): @overload def __ge__(self, other: Self) -> bool: ... @overload - def __ge__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... + def __ge__(self, other: PeriodIndex) -> np_1darray_bool: ... @overload def __ge__( self, other: Series[Period] # pyrefly: ignore[bad-specialization] @@ -129,7 +129,7 @@ class Period(PeriodMixin): @overload def __gt__(self, other: Self) -> bool: ... @overload - def __gt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... + def __gt__(self, other: PeriodIndex) -> np_1darray_bool: ... @overload def __gt__( self, other: Series[Period] # pyrefly: ignore[bad-specialization] @@ -141,7 +141,7 @@ class Period(PeriodMixin): @overload def __le__(self, other: Self) -> bool: ... @overload - def __le__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... + def __le__(self, other: PeriodIndex) -> np_1darray_bool: ... @overload def __le__( self, other: Series[Period] # pyrefly: ignore[bad-specialization] @@ -153,7 +153,7 @@ class Period(PeriodMixin): @overload def __lt__(self, other: Self) -> bool: ... @overload - def __lt__(self, other: PeriodIndex) -> np_1darray[np.bool]: ... + def __lt__(self, other: PeriodIndex) -> np_1darray_bool: ... @overload def __lt__( self, other: Series[Period] # pyrefly: ignore[bad-specialization] @@ -167,7 +167,7 @@ class Period(PeriodMixin): @overload def __ne__(self, other: Self) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload - def __ne__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __ne__(self, other: Index) -> np_1darray_bool: ... # type: ignore[overload-overlap] @overload def __ne__(self, other: Series[Period]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 308af041e..5d808bb09 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -29,7 +29,7 @@ from pandas._typing import ( Just, ShapeT, TimeUnit, - np_1darray, + np_1darray_bool, np_ndarray, ) @@ -274,7 +274,7 @@ class Timedelta(timedelta): @overload def __eq__(self, other: Series[Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __eq__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __eq__(self, other: Index) -> np_1darray_bool: ... # type: ignore[overload-overlap] @overload def __eq__( # type: ignore[overload-overlap] self, other: np_ndarray[ShapeT, np.timedelta64] @@ -287,7 +287,7 @@ class Timedelta(timedelta): @overload def __ne__(self, other: Series[Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __ne__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __ne__(self, other: Index) -> np_1darray_bool: ... # type: ignore[overload-overlap] @overload def __ne__( # type: ignore[overload-overlap] self, other: np_ndarray[ShapeT, np.timedelta64] @@ -316,7 +316,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __le__(self, other: timedelta | np.timedelta64 | Self) -> bool: ... @overload - def __le__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + def __le__(self, other: TimedeltaIndex) -> np_1darray_bool: ... @overload def __le__( self, other: np_ndarray[ShapeT, np.timedelta64] @@ -327,7 +327,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __lt__(self, other: timedelta | np.timedelta64 | Self) -> bool: ... @overload - def __lt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + def __lt__(self, other: TimedeltaIndex) -> np_1darray_bool: ... @overload def __lt__( self, other: np_ndarray[ShapeT, np.timedelta64] @@ -338,7 +338,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __ge__(self, other: timedelta | np.timedelta64 | Self) -> bool: ... @overload - def __ge__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + def __ge__(self, other: TimedeltaIndex) -> np_1darray_bool: ... @overload def __ge__( self, other: np_ndarray[ShapeT, np.timedelta64] @@ -349,7 +349,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __gt__(self, other: timedelta | np.timedelta64 | Self) -> bool: ... @overload - def __gt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... + def __gt__(self, other: TimedeltaIndex) -> np_1darray_bool: ... @overload def __gt__( self, other: np_ndarray[ShapeT, np.timedelta64] diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index a2b329338..5dbb28c8f 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -37,7 +37,7 @@ from pandas._typing import ( ShapeT, TimestampNonexistent, TimeUnit, - np_1darray, + np_1darray_bool, np_ndarray, ) @@ -174,7 +174,7 @@ class Timestamp(datetime, SupportsIndex): @overload # type: ignore[override] def __le__(self, other: datetime | np.datetime64 | Self) -> bool: ... @overload - def __le__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + def __le__(self, other: DatetimeIndex) -> np_1darray_bool: ... @overload def __le__( self, other: np_ndarray[ShapeT, np.datetime64] @@ -184,7 +184,7 @@ class Timestamp(datetime, SupportsIndex): @overload # type: ignore[override] def __lt__(self, other: datetime | np.datetime64 | Self) -> bool: ... @overload - def __lt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + def __lt__(self, other: DatetimeIndex) -> np_1darray_bool: ... @overload def __lt__( self, other: np_ndarray[ShapeT, np.datetime64] @@ -194,7 +194,7 @@ class Timestamp(datetime, SupportsIndex): @overload # type: ignore[override] def __ge__(self, other: datetime | np.datetime64 | Self) -> bool: ... @overload - def __ge__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + def __ge__(self, other: DatetimeIndex) -> np_1darray_bool: ... @overload def __ge__( self, other: np_ndarray[ShapeT, np.datetime64] @@ -204,7 +204,7 @@ class Timestamp(datetime, SupportsIndex): @overload # type: ignore[override] def __gt__(self, other: datetime | np.datetime64 | Self) -> bool: ... @overload - def __gt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ... + def __gt__(self, other: DatetimeIndex) -> np_1darray_bool: ... @overload def __gt__( self, other: np_ndarray[ShapeT, np.datetime64] @@ -244,7 +244,7 @@ class Timestamp(datetime, SupportsIndex): @overload def __eq__(self, other: Series[Timestamp]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __eq__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __eq__(self, other: Index) -> np_1darray_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 @@ -254,7 +254,7 @@ class Timestamp(datetime, SupportsIndex): @overload def __ne__(self, other: Series[Timestamp]) -> Series[bool]: ... # type: ignore[overload-overlap] @overload - def __ne__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] + def __ne__(self, other: Index) -> np_1darray_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 diff --git a/pandas-stubs/_testing/__init__.pyi b/pandas-stubs/_testing/__init__.pyi index 3db4eccbb..57d12c7e7 100644 --- a/pandas-stubs/_testing/__init__.pyi +++ b/pandas-stubs/_testing/__init__.pyi @@ -12,7 +12,6 @@ from typing import ( import warnings from matplotlib.artist import Artist -import numpy as np from pandas import ( Categorical, DataFrame, @@ -31,6 +30,7 @@ from pandas.core.arrays.base import ExtensionArray from pandas._typing import ( AnyArrayLike, T, + np_ndarray, ) def assert_almost_equal( @@ -61,7 +61,7 @@ def assert_attr_equal( attr: str, left: object, right: object, obj: str = "Attributes" ) -> None: ... def assert_is_valid_plot_return_object( - objs: Series | np.ndarray | Artist | tuple | dict, + objs: Series | np_ndarray | Artist | tuple | dict, ) -> None: ... def assert_is_sorted(seq: AnyArrayLike) -> None: ... def assert_categorical_equal( @@ -96,7 +96,7 @@ def assert_extension_array_equal( left: ExtensionArray, right: ExtensionArray, check_dtype: bool | Literal["equiv"] = True, - index_values: Index | np.ndarray | None = None, + index_values: Index | np_ndarray | None = None, check_exact: bool = False, rtol: float = 1e-5, atol: float = 1e-8, diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index ef1398965..6cb6bfa92 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -85,7 +85,7 @@ HashableT5 = TypeVar("HashableT5", bound=Hashable) # array-like -ArrayLike: TypeAlias = ExtensionArray | np.ndarray +ArrayLike: TypeAlias = ExtensionArray | npt.NDArray[Any] AnyArrayLike: TypeAlias = ArrayLike | Index | Series AnyArrayLikeInt: TypeAlias = ( IntegerArray | Index[int] | Series[int] | npt.NDArray[np.integer] @@ -166,6 +166,7 @@ ToTimestampHow: TypeAlias = Literal["s", "e", "start", "end"] NDFrameT = TypeVar("NDFrameT", bound=NDFrame) IndexT = TypeVar("IndexT", bound=Index) +T_EXTENSION_ARRAY = TypeVar("T_EXTENSION_ARRAY", bound=ExtensionArray) # From _typing.py, not used here: # FreqIndexT = TypeVar("FreqIndexT", "DatetimeIndex", "PeriodIndex", "TimedeltaIndex") @@ -697,11 +698,11 @@ InterpolateOptions: TypeAlias = Literal[ # Using List[int] here rather than Sequence[int] to disallow tuples. ScalarIndexer: TypeAlias = int | np.integer -SequenceIndexer: TypeAlias = slice | list[int] | np.ndarray +SequenceIndexer: TypeAlias = slice | list[int] | npt.NDArray[np.integer | np.bool] PositionalIndexer: TypeAlias = ScalarIndexer | SequenceIndexer PositionalIndexerTuple: TypeAlias = tuple[PositionalIndexer, PositionalIndexer] # PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple] Not used in stubs -TakeIndexer: TypeAlias = Sequence[int] | Sequence[np.integer] | npt.NDArray[np.integer] +TakeIndexer: TypeAlias = Sequence[int | np.integer] | npt.NDArray[np.integer | np.bool] # Shared by functions such as drop and astype IgnoreRaise: TypeAlias = Literal["ignore", "raise"] @@ -783,7 +784,7 @@ TimeNonexistent: TypeAlias = ( DropKeep: TypeAlias = Literal["first", "last", False] CorrelationMethod: TypeAlias = ( Literal["pearson", "kendall", "spearman"] - | Callable[[np.ndarray, np.ndarray], float] + | Callable[[np.typing.NDArray[Any], np.typing.NDArray[Any]], float] ) AlignJoin: TypeAlias = Literal["outer", "inner", "left", "right"] DtypeBackend: TypeAlias = Literal["pyarrow", "numpy_nullable"] @@ -832,31 +833,18 @@ SliceType: TypeAlias = Hashable | None ## All types below this point are only used in pandas-stubs ###### -num: TypeAlias = complex - -DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) -KeysArgType: TypeAlias = Any -ListLikeT = TypeVar("ListLikeT", bound=ListLike) -ListLikeExceptSeriesAndStr: TypeAlias = ( - MutableSequence[Any] | np.ndarray | tuple[Any, ...] | Index -) -ListLikeU: TypeAlias = Sequence | np.ndarray | Series | Index -ListLikeHashable: TypeAlias = ( - MutableSequence[HashableT] | np.ndarray | tuple[HashableT, ...] | range -) StrLike: TypeAlias = str | np.str_ ScalarT = TypeVar("ScalarT", bound=Scalar) # Refine the definitions below in 3.9 to use the specialized type. +np_num: TypeAlias = np.bool | np.integer | np.floating | np.complexfloating +np_ndarray_intp: TypeAlias = npt.NDArray[np.intp] np_ndarray_int64: TypeAlias = npt.NDArray[np.int64] np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger] np_ndarray_anyint: TypeAlias = npt.NDArray[np.integer] np_ndarray_float: TypeAlias = npt.NDArray[np.floating] np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating] np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_] -np_ndarray_num: TypeAlias = npt.NDArray[ - np.bool | np.integer | np.floating | np.complexfloating -] np_ndarray_str: TypeAlias = npt.NDArray[np.str_] np_ndarray_dt: TypeAlias = npt.NDArray[np.datetime64] np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64] @@ -867,13 +855,35 @@ GenericT_co = TypeVar("GenericT_co", bound=np.generic, default=Any, covariant=Tr GenericT_contra = TypeVar( "GenericT_contra", bound=np.generic, default=Any, contravariant=True ) +NpNumT = TypeVar("NpNumT", bound=np_num, default=np_num) ShapeT = TypeVar("ShapeT", bound=tuple[int, ...], default=tuple[Any, ...]) # Numpy ndarray with more ergonomic typevar np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]] +np_ndarray_num: TypeAlias = np_ndarray[ShapeT, NpNumT] + # 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_1darray_bool: TypeAlias = np_1darray[np.bool] +np_1darray_intp: TypeAlias = np_1darray[np.intp] +np_1darray_int64: TypeAlias = np_1darray[np.int64] +np_1darray_anyint: TypeAlias = np_1darray[np.integer] +np_1darray_float: TypeAlias = np_1darray[np.floating] +np_1darray_dt: TypeAlias = np_1darray[np.datetime64] +np_1darray_td: TypeAlias = np_1darray[np.timedelta64] + np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] +DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) +KeysArgType: TypeAlias = Any +ListLikeT = TypeVar("ListLikeT", bound=ListLike) +ListLikeExceptSeriesAndStr: TypeAlias = ( + MutableSequence[Any] | np_1darray | tuple[Any, ...] | Index +) +ListLikeU: TypeAlias = Sequence | np_1darray | Series | Index +ListLikeHashable: TypeAlias = ( + MutableSequence[HashableT] | np_1darray | tuple[HashableT, ...] | range +) + class SupportsDType(Protocol[GenericT_co]): @property def dtype(self) -> np.dtype[GenericT_co]: ... diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 5174218c6..a529e733f 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -4,66 +4,90 @@ from typing import ( overload, ) -import numpy as np -from pandas import ( - Categorical, - CategoricalIndex, - Index, - IntervalIndex, - PeriodIndex, - Series, -) +from numpy import typing as npt from pandas.api.extensions import ExtensionArray +from pandas.core.arrays.categorical import Categorical +from pandas.core.indexes.base import Index +from pandas.core.indexes.category import CategoricalIndex +from pandas.core.indexes.datetimes import DatetimeIndex +from pandas.core.indexes.interval import IntervalIndex +from pandas.core.indexes.multi import MultiIndex +from pandas.core.indexes.period import PeriodIndex +from pandas.core.indexes.range import RangeIndex +from pandas.core.indexes.timedeltas import TimedeltaIndex +from pandas.core.series import Series from pandas._typing import ( + T_EXTENSION_ARRAY, AnyArrayLike, + GenericT, IntervalT, TakeIndexer, np_1darray, + np_1darray_dt, + np_1darray_int64, + np_1darray_td, + np_ndarray, ) # These are type: ignored because the Index types overlap due to inheritance but indices # with extension types return the same type while standard type return ndarray - -@overload -def unique( # pyright: ignore[reportOverlappingOverload] - values: PeriodIndex, -) -> PeriodIndex: ... @overload -def unique(values: CategoricalIndex) -> CategoricalIndex: ... # type: ignore[overload-overlap] +def unique(values: CategoricalIndex) -> CategoricalIndex: ... @overload def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ... @overload -def unique(values: Index) -> np.ndarray: ... +def unique(values: PeriodIndex) -> PeriodIndex: ... +@overload +# switch to DatetimeIndex after Pandas 3.0 +def unique(values: DatetimeIndex) -> np_1darray_dt | DatetimeIndex: ... +@overload +# switch to DatetimeIndex after Pandas 3.0 +def unique(values: TimedeltaIndex) -> np_1darray_td: ... +@overload +# switch to Index[int] after Pandas 3.0 +def unique(values: RangeIndex) -> np_1darray_int64: ... +@overload +def unique(values: MultiIndex) -> np_ndarray: ... +@overload +def unique(values: Index) -> np_1darray | Index: ... # switch to Index after Pandas 3.0 @overload def unique(values: Categorical) -> Categorical: ... + +# @overload +# def unique(values: Series[Never]) -> np_1darray | ExtensionArray: ... +# TODO: DatetimeArray python/mypy#19952 +# @overload +# def unique(values: Series[Timestamp]) -> np_ndarray_dt | ExtensionArray: ... +# @overload +# def unique(values: Series[int]) -> np_1darray_anyint | ExtensionArray: ... @overload -def unique(values: Series) -> np.ndarray | ExtensionArray: ... +def unique(values: Series) -> np_1darray | ExtensionArray: ... @overload -def unique(values: np.ndarray) -> np.ndarray: ... +def unique(values: npt.NDArray[GenericT]) -> np_1darray[GenericT]: ... @overload -def unique(values: ExtensionArray) -> ExtensionArray: ... +def unique(values: T_EXTENSION_ARRAY) -> T_EXTENSION_ARRAY: ... @overload def factorize( - values: np.ndarray, + values: npt.NDArray[GenericT], sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np.ndarray, np.ndarray]: ... +) -> tuple[np_1darray_int64, np_1darray[GenericT]]: ... @overload def factorize( values: Index | Series, sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np_1darray, Index]: ... +) -> tuple[np_1darray_int64, Index]: ... @overload def factorize( values: Categorical, sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np_1darray, Categorical]: ... +) -> tuple[np_1darray_int64, Categorical]: ... def value_counts( values: AnyArrayLike | list | tuple, sort: bool = True, @@ -73,7 +97,7 @@ def value_counts( dropna: bool = True, ) -> Series: ... def take( - arr: np.ndarray | ExtensionArray | Index | Series, + arr: np_ndarray[Any] | ExtensionArray | Index | Series, indices: TakeIndexer, axis: Literal[0, 1] = 0, allow_fill: bool = False, diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index 09bf2b0cb..7ce4f26c6 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -22,6 +22,8 @@ from pandas._typing import ( SequenceIndexer, TakeIndexer, np_1darray, + np_1darray_intp, + np_ndarray, npt, ) @@ -35,7 +37,7 @@ class ExtensionArray: def __getitem__(self, item: ScalarIndexer) -> Any: ... @overload def __getitem__(self, item: SequenceIndexer) -> Self: ... - def __setitem__(self, key: int | slice | np.ndarray, value: Any) -> None: ... + def __setitem__(self, key: int | slice | np_ndarray, value: Any) -> None: ... def __len__(self) -> int: ... def __iter__(self) -> Iterator[Any]: ... def __contains__(self, item: object) -> bool | np.bool_: ... @@ -44,7 +46,7 @@ class ExtensionArray: dtype: npt.DTypeLike | None = ..., copy: bool = False, na_value: Scalar = ..., - ) -> np_1darray[Any]: ... + ) -> np_1darray: ... @property def dtype(self) -> ExtensionDtype: ... @property @@ -75,7 +77,7 @@ class ExtensionArray: value: ListLike, side: Literal["left", "right"] = ..., sorter: ListLike | None = ..., - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... @overload def searchsorted( self, diff --git a/pandas-stubs/core/arrays/boolean.pyi b/pandas-stubs/core/arrays/boolean.pyi index 224bae066..ed0b61477 100644 --- a/pandas-stubs/core/arrays/boolean.pyi +++ b/pandas-stubs/core/arrays/boolean.pyi @@ -1,10 +1,12 @@ from typing import Any -import numpy as np from pandas.core.arrays.masked import BaseMaskedArray as BaseMaskedArray from pandas._libs.missing import NAType -from pandas._typing import type_t +from pandas._typing import ( + np_ndarray_bool, + type_t, +) from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype @@ -16,7 +18,7 @@ class BooleanDtype(ExtensionDtype): class BooleanArray(BaseMaskedArray): def __init__( - self, values: np.ndarray, mask: np.ndarray, copy: bool = ... + self, values: np_ndarray_bool, mask: np_ndarray_bool, copy: bool = ... ) -> None: ... @property def dtype(self): ... diff --git a/pandas-stubs/core/arrays/categorical.pyi b/pandas-stubs/core/arrays/categorical.pyi index 667d46819..3ea880995 100644 --- a/pandas-stubs/core/arrays/categorical.pyi +++ b/pandas-stubs/core/arrays/categorical.pyi @@ -25,6 +25,7 @@ from pandas._typing import ( SequenceIndexer, TakeIndexer, np_1darray, + np_1darray_bool, ) from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype @@ -95,10 +96,10 @@ class Categorical(ExtensionArray): @property def nbytes(self) -> int: ... def memory_usage(self, deep: 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 isna(self) -> np_1darray_bool: ... + def isnull(self) -> np_1darray_bool: ... + def notna(self) -> np_1darray_bool: ... + def notnull(self) -> np_1darray_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/datetimelike.pyi b/pandas-stubs/core/arrays/datetimelike.pyi index 1f4870898..f392f64f2 100644 --- a/pandas-stubs/core/arrays/datetimelike.pyi +++ b/pandas-stubs/core/arrays/datetimelike.pyi @@ -5,7 +5,6 @@ from typing import ( overload, ) -import numpy as np from pandas.core.arrays.base import ( ExtensionArray, ExtensionOpsMixin, @@ -65,8 +64,6 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, ExtensionArray): def ravel(self, *args: Any, **kwargs: Any): ... def __iter__(self): ... @property - def asi8(self) -> np.ndarray: ... - @property def nbytes(self): ... def __array__( self, dtype: NpDtype | None = None, copy: bool | None = None diff --git a/pandas-stubs/core/arrays/interval.pyi b/pandas-stubs/core/arrays/interval.pyi index 52d23ac51..cdad76567 100644 --- a/pandas-stubs/core/arrays/interval.pyi +++ b/pandas-stubs/core/arrays/interval.pyi @@ -4,7 +4,6 @@ from typing import ( overload, ) -import numpy as np from pandas.core.arrays.base import ExtensionArray as ExtensionArray from pandas.core.indexes.base import Index from pandas.core.series import Series @@ -22,6 +21,8 @@ from pandas._typing import ( SequenceIndexer, TakeIndexer, np_1darray, + np_1darray_bool, + np_ndarray, ) IntervalOrNA: TypeAlias = Interval | float @@ -105,6 +106,6 @@ class IntervalArray(IntervalMixin, ExtensionArray): def contains(self, other: Series) -> Series[bool]: ... @overload def contains( - self, other: Scalar | ExtensionArray | Index | np.ndarray - ) -> np_1darray[np.bool]: ... + self, other: Scalar | ExtensionArray | Index | np_ndarray + ) -> np_1darray_bool: ... def overlaps(self, other: Interval) -> bool: ... diff --git a/pandas-stubs/core/arrays/masked.pyi b/pandas-stubs/core/arrays/masked.pyi index c382e8e67..89bf39012 100644 --- a/pandas-stubs/core/arrays/masked.pyi +++ b/pandas-stubs/core/arrays/masked.pyi @@ -3,7 +3,6 @@ from typing import ( overload, ) -import numpy as np from pandas.core.arrays import ( ExtensionArray as ExtensionArray, ExtensionOpsMixin, @@ -31,7 +30,7 @@ class BaseMaskedArray(ExtensionArray, ExtensionOpsMixin): dtype: npt.DTypeLike | None = ..., copy: bool = False, na_value: Scalar = ..., - ) -> np.ndarray: ... + ) -> np_1darray: ... __array_priority__: int = ... def __array__( self, dtype: NpDtype | None = None, copy: bool | None = None diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index 826805678..1e3a93bf9 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -43,6 +43,8 @@ from pandas._typing import ( Scalar, SupportsDType, np_1darray, + np_1darray_intp, + np_ndarray, np_ndarray_anyint, np_ndarray_bool, np_ndarray_complex, @@ -162,7 +164,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]): value: ListLike, side: Literal["left", "right"] = ..., sorter: ListLike | None = None, - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... @overload def searchsorted( self, @@ -255,7 +257,7 @@ NumListLike: TypeAlias = ( # TODO: pandas-dev/pandas-stubs#1474 deprecated, do | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex - | dict[str, np.ndarray] + | dict[str, np_ndarray] | Sequence[complex] ) diff --git a/pandas-stubs/core/computation/eval.pyi b/pandas-stubs/core/computation/eval.pyi index 19df29909..b1bb890e4 100644 --- a/pandas-stubs/core/computation/eval.pyi +++ b/pandas-stubs/core/computation/eval.pyi @@ -12,7 +12,7 @@ from pandas.core.computation.ops import BinOp from pandas._typing import ( Scalar, - npt, + np_ndarray, ) def eval( @@ -25,4 +25,4 @@ def eval( level: int = 0, target: object | None = None, inplace: bool = False, -) -> npt.NDArray | Scalar | DataFrame | Series | None: ... +) -> np_ndarray | Scalar | DataFrame | Series | None: ... diff --git a/pandas-stubs/core/construction.pyi b/pandas-stubs/core/construction.pyi index 2372c1edc..58872d197 100644 --- a/pandas-stubs/core/construction.pyi +++ b/pandas-stubs/core/construction.pyi @@ -3,16 +3,28 @@ from typing import overload import numpy as np from pandas.core.arrays.base import ExtensionArray +from pandas.core.arrays.boolean import BooleanArray from pandas.core.arrays.integer import IntegerArray from pandas._libs.missing import NAType +from pandas._typing import ( + BooleanDtypeArg, + IntDtypeArg, + UIntDtypeArg, +) from pandas.core.dtypes.dtypes import ExtensionDtype +@overload +def array( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + data: Sequence[bool | NAType | None], + dtype: BooleanDtypeArg | None = None, + copy: bool = True, +) -> BooleanArray: ... @overload def array( - data: Sequence[int] | Sequence[int | NAType], - dtype: str | np.dtype | ExtensionDtype | None = None, + data: Sequence[int | NAType | None], + dtype: IntDtypeArg | UIntDtypeArg | None = None, copy: bool = True, ) -> IntegerArray: ... @overload diff --git a/pandas-stubs/core/dtypes/missing.pyi b/pandas-stubs/core/dtypes/missing.pyi index d53c0cf53..dad3e5c2e 100644 --- a/pandas-stubs/core/dtypes/missing.pyi +++ b/pandas-stubs/core/dtypes/missing.pyi @@ -18,7 +18,7 @@ from pandas._typing import ( Scalar, ScalarT, ShapeT, - np_1darray, + np_1darray_bool, np_ndarray, np_ndarray_bool, ) @@ -31,7 +31,7 @@ def isna(obj: DataFrame) -> DataFrame: ... @overload def isna(obj: Series) -> Series[bool]: ... @overload -def isna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray[np.bool]: ... +def isna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray_bool: ... @overload def isna(obj: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ... @overload @@ -48,7 +48,7 @@ def notna(obj: DataFrame) -> DataFrame: ... @overload def notna(obj: Series) -> Series[bool]: ... @overload -def notna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray[np.bool]: ... +def notna(obj: Index | ExtensionArray | list[ScalarT]) -> np_1darray_bool: ... @overload def notna(obj: np_ndarray[ShapeT]) -> np_ndarray[ShapeT, np.bool]: ... @overload diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 89fc4fc0d..59a0937a2 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -31,6 +31,7 @@ from typing import ( from matplotlib.axes import Axes as PlotAxes import numpy as np +from numpy import typing as npt from pandas import ( Period, Timedelta, @@ -163,8 +164,10 @@ from pandas._typing import ( WriteBuffer, XMLParsers, np_2darray, - npt, - num, + np_ndarray, + np_ndarray_bool, + np_ndarray_float, + np_ndarray_num, ) from pandas.io.formats.style import Styler @@ -208,7 +211,7 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): Scalar | Series | DataFrame - | np.ndarray + | np_ndarray | NAType | NaTType | Mapping[Hashable, Scalar | NAType | NaTType] @@ -426,8 +429,8 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def __matmul__(self, other: Series) -> Series: ... @overload - def __matmul__(self, other: np.ndarray) -> Self: ... - def __rmatmul__(self, other): ... + def __matmul__(self, other: np_ndarray_num) -> Self: ... + def __rmatmul__(self, other: np_ndarray_num) -> Self: ... @overload @classmethod def from_dict( @@ -831,7 +834,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def eval( self, expr: _str, *, inplace: Literal[False] = False, **kwargs: Any - ) -> Scalar | np.ndarray | Self | Series: ... + ) -> Scalar | np_ndarray | Self | Series: ... @overload def select_dtypes( self, include: StrDtypeArg, exclude: _AstypeArgExtList | None = ... @@ -969,7 +972,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): Label | Series | Index - | np.ndarray + | np_ndarray | Iterator[Hashable] | Sequence[Hashable] ), @@ -986,7 +989,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): Label | Series | Index - | np.ndarray + | np_ndarray | Iterator[Hashable] | Sequence[Hashable] ), @@ -1403,8 +1406,8 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): ) -> Self | Series: ... def melt( self, - id_vars: tuple | Sequence | np.ndarray | None = ..., - value_vars: tuple | Sequence | np.ndarray | None = ..., + id_vars: tuple | Sequence | np_ndarray | None = ..., + value_vars: tuple | Sequence | np_ndarray | None = ..., var_name: Scalar | None = None, value_name: Scalar = "value", col_level: int | _str | None = ..., @@ -1677,7 +1680,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def quantile( self, - q: list[float] | np.ndarray, + q: Sequence[float] | np_ndarray_float, axis: Axis = ..., numeric_only: _bool = ..., interpolation: QuantileInterpolation = ..., @@ -1844,7 +1847,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def __add__(self, other: Any) -> Self: ... def add( self, - other: num | ListLike | DataFrame, + other: complex | ListLike | DataFrame, axis: Axis | None = "columns", level: Level | None = None, fill_value: float | None = None, @@ -1852,7 +1855,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def __radd__(self, other: Any) -> Self: ... def radd( self, - other, + other: complex | ListLike | DataFrame, axis: Axis = "columns", level: Level | None = None, fill_value: float | None = None, @@ -1860,7 +1863,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def __sub__(self, other: Any) -> Self: ... def sub( self, - other: num | ListLike | DataFrame, + other: complex | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, @@ -1868,7 +1871,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def __rsub__(self, other: Any) -> Self: ... def rsub( self, - other, + other: complex | ListLike | DataFrame, axis: Axis = ..., level: Level | None = ..., fill_value: float | None = None, @@ -2071,20 +2074,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): include: Literal["all"] | list[Dtype] | None = ..., exclude: list[Dtype] | None = ..., ) -> Self: ... - def div( - self, - other: num | ListLike | DataFrame, - axis: Axis | None = ..., - level: Level | None = ..., - fill_value: float | None = None, - ) -> Self: ... - def divide( + def truediv( self, - other: num | ListLike | DataFrame, + other: complex | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, ) -> Self: ... + div = truediv @final def droplevel(self, level: Level | list[Level], axis: Axis = 0) -> Self: ... def eq(self, other, axis: Axis = "columns", level: Level | None = ...) -> Self: ... @@ -2138,7 +2135,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def first_valid_index(self) -> Scalar: ... def floordiv( self, - other: num | ListLike | DataFrame, + other: float | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, @@ -2206,7 +2203,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | np_ndarray_bool | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2222,7 +2219,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | np_ndarray_bool | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2262,21 +2259,14 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): ) -> Series: ... def mod( self, - other: num | ListLike | DataFrame, + other: float | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, ) -> Self: ... def mul( self, - other: num | ListLike | DataFrame, - axis: Axis | None = ..., - level: Level | None = ..., - fill_value: float | None = None, - ) -> Self: ... - def multiply( - self, - other: num | ListLike | DataFrame, + other: complex | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, @@ -2294,7 +2284,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def pop(self, item: _str) -> Series: ... def pow( self, - other: num | ListLike | DataFrame, + other: complex | ListLike | DataFrame, axis: Axis | None = ..., level: Level | None = ..., fill_value: float | None = None, @@ -2616,13 +2606,6 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): ) -> _str: ... @final def to_xarray(self) -> xr.Dataset: ... - def truediv( - self, - other: num | ListLike | DataFrame, - axis: Axis | None = "columns", - level: Level | None = None, - fill_value: float | None = None, - ) -> Self: ... @final def truncate( self, @@ -2663,7 +2646,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | np_ndarray_bool | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2679,7 +2662,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | np_ndarray_bool | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 1b292b843..d5bee980b 100644 --- a/pandas-stubs/core/generic.pyi +++ b/pandas-stubs/core/generic.pyi @@ -21,7 +21,6 @@ from typing import ( overload, ) -import numpy as np from pandas import Index import pandas.core.indexing as indexing from pandas.core.resample import DatetimeIndexResampler @@ -58,6 +57,7 @@ from pandas._typing import ( TimeGrouperOrigin, TimestampConvertibleTypes, WriteBuffer, + np_ndarray, ) from pandas.io.pytables import HDFStore @@ -100,7 +100,7 @@ class NDFrame(indexing.IndexingMixin): @property def empty(self) -> _bool: ... __array_priority__: int = ... - def __array__(self, dtype=...) -> np.ndarray: ... + def __array__(self, dtype=...) -> np_ndarray: ... @final def to_excel( self, diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 7c058cd79..bbd764d77 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -19,7 +19,6 @@ from typing import ( ) from matplotlib.axes import Axes as PlotAxes -import numpy as np from pandas.core.frame import DataFrame from pandas.core.groupby.base import TransformReductionListType from pandas.core.groupby.groupby import ( @@ -47,6 +46,7 @@ from pandas._typing import ( TakeIndexer, WindowingEngine, WindowingEngineKwargs, + np_ndarray, ) AggScalar: TypeAlias = str | Callable[..., Any] @@ -418,7 +418,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]): def plot(self) -> GroupByPlot[Self]: ... def corr( self, - method: str | Callable[[np.ndarray, np.ndarray], float] = ..., + method: str | Callable[[np_ndarray, np_ndarray], float] = ..., min_periods: int = ..., numeric_only: bool = False, ) -> DataFrame: ... diff --git a/pandas-stubs/core/groupby/groupby.pyi b/pandas-stubs/core/groupby/groupby.pyi index c5bc50110..b60864ada 100644 --- a/pandas-stubs/core/groupby/groupby.pyi +++ b/pandas-stubs/core/groupby/groupby.pyi @@ -64,7 +64,8 @@ from pandas._typing import ( TimestampConvertibleTypes, WindowingEngine, WindowingEngineKwargs, - npt, + np_ndarray_dt, + np_ndarray_int64, ) from pandas.plotting import PlotAccessor @@ -240,7 +241,7 @@ class GroupBy(BaseGroupBy[NDFrameT]): adjust: bool = ..., ignore_na: bool = ..., axis: Axis = ..., - times: str | np.ndarray | Series | np.timedelta64 | None = ..., + times: str | np_ndarray_dt | Series | np.timedelta64 | None = ..., method: CalculationMethod = ..., *, selection: IndexLabel | None = ..., @@ -357,7 +358,7 @@ class BaseGroupBy(SelectionMixin[NDFrameT], GroupByIndexingMixin): def ngroups(self) -> int: ... @final @property - def indices(self) -> dict[Hashable, Index | npt.NDArray[np.int_] | list[int]]: ... + def indices(self) -> dict[Hashable, Index | np_ndarray_int64 | list[int]]: ... @overload def pipe( self, diff --git a/pandas-stubs/core/groupby/grouper.pyi b/pandas-stubs/core/groupby/grouper.pyi index 8ac7f901a..b26b342b5 100644 --- a/pandas-stubs/core/groupby/grouper.pyi +++ b/pandas-stubs/core/groupby/grouper.pyi @@ -1,34 +1,21 @@ -from collections.abc import ( - Hashable, - Iterator, -) from typing import ( Any, final, overload, ) -import numpy as np -from pandas import ( - DataFrame, - Index, - Series, -) from pandas.core.resample import TimeGrouper from typing_extensions import Self from pandas._libs.lib import _NoDefaultDoNotUse from pandas._typing import ( - ArrayLike, Axis, Frequency, Incomplete, KeysArgType, Level, ListLikeHashable, - npt, ) -from pandas.util._decorators import cache_readonly class Grouper: key: KeysArgType | None @@ -51,25 +38,3 @@ class Grouper: def __new__(cls, *args: Any, freq: Frequency, **kwargs: Any) -> TimeGrouper: ... @final def __repr__(self) -> str: ... # noqa: PYI029 __repr__ here is final - -@final -class Grouping: - level: Level | None - obj: DataFrame | Series | None - in_axis: bool - grouping_vector: Incomplete - def __iter__(self) -> Iterator[Hashable]: ... - @cache_readonly - def name(self) -> Hashable: ... - @cache_readonly - def ngroups(self) -> int: ... - @cache_readonly - def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]: ... - @property - def codes(self) -> npt.NDArray[np.signedinteger]: ... - @cache_readonly - def group_arraylike(self) -> ArrayLike: ... - @cache_readonly - def result_index(self) -> Index: ... - @cache_readonly - def group_index(self) -> Index: ... diff --git a/pandas-stubs/core/groupby/ops.pyi b/pandas-stubs/core/groupby/ops.pyi index 740eeb760..df5cb379a 100644 --- a/pandas-stubs/core/groupby/ops.pyi +++ b/pandas-stubs/core/groupby/ops.pyi @@ -1,96 +1,17 @@ from collections.abc import ( - Callable, - Hashable, Iterator, ) from typing import ( Generic, - final, ) import numpy as np -from pandas import ( - Index, - Series, -) -from pandas.core.groupby import grouper from pandas._typing import ( - ArrayLike, AxisInt, - Incomplete, NDFrameT, - Shape, - T, npt, ) -from pandas.util._decorators import cache_readonly - -class BaseGrouper: - axis: Index - dropna: bool - @property - def groupings(self) -> list[grouper.Grouping]: ... - @property - def shape(self) -> Shape: ... - def __iter__(self) -> Iterator: ... - @property - def nkeys(self) -> int: ... - def get_iterator( - self, data: NDFrameT, axis: AxisInt = ... - ) -> Iterator[tuple[Hashable, NDFrameT]]: ... - @cache_readonly - def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]: ... - @final - def result_ilocs(self) -> npt.NDArray[np.intp]: ... - @final - @property - def codes(self) -> list[npt.NDArray[np.signedinteger]]: ... - @property - def levels(self) -> list[Index]: ... - @property - def names(self) -> list: ... - @final - def size(self) -> Series: ... - @cache_readonly - def groups(self) -> dict[Hashable, np.ndarray]: ... - @final - @cache_readonly - def is_monotonic(self) -> bool: ... - @final - @cache_readonly - def has_dropped_na(self) -> bool: ... - @cache_readonly - def group_info(self) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], int]: ... - @cache_readonly - def codes_info(self) -> npt.NDArray[np.intp]: ... - @final - @cache_readonly - def ngroups(self) -> int: ... - @property - def reconstructed_codes(self) -> list[npt.NDArray[np.intp]]: ... - @cache_readonly - def result_index(self) -> Index: ... - @final - def get_group_levels(self) -> list[ArrayLike]: ... - @final - def agg_series( - self, - obj: Series, - func: Callable[[Series], object], - preserve_dtype: bool = ..., - ) -> ArrayLike: ... - @final - def apply_groupwise( - self, f: Callable[[NDFrameT], T], data: NDFrameT, axis: AxisInt = ... - ) -> tuple[list[T], bool]: ... - -class BinGrouper(BaseGrouper): - bins: npt.NDArray[np.int64] - binlabels: Index - indexer: npt.NDArray[np.intp] - @cache_readonly - def indices(self) -> dict[Incomplete, list[int]]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] class DataSplitter(Generic[NDFrameT]): data: NDFrameT diff --git a/pandas-stubs/core/indexers.pyi b/pandas-stubs/core/indexers.pyi deleted file mode 100644 index c6bcc6d84..000000000 --- a/pandas-stubs/core/indexers.pyi +++ /dev/null @@ -1,46 +0,0 @@ -from typing import Any - -import numpy as np - -def check_array_indexer(arrayArrayLike, indexer): ... - -class BaseIndexer: - def __init__( - self, - index_array: np.ndarray | None = ..., - window_size: int = ..., - **kwargs: Any, - ) -> None: ... - def get_window_bounds( - self, - num_values: int = ..., - min_periods: int | None = ..., - center: bool | None = ..., - closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... - -class VariableOffsetWindowIndexer(BaseIndexer): - def __init__( - self, - index_array: np.ndarray | None = ..., - window_size: int = ..., - index=..., - offset=..., - **kwargs: Any, - ) -> None: ... - def get_window_bounds( - self, - num_values: int = ..., - min_periods: int | None = ..., - center: bool | None = ..., - closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... - -class FixedForwardWindowIndexer(BaseIndexer): - def get_window_bounds( - self, - num_values: int = ..., - min_periods: int | None = ..., - center: bool | None = ..., - closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... diff --git a/pandas-stubs/core/indexers/objects.pyi b/pandas-stubs/core/indexers/objects.pyi index 2187cf885..1a5cdd9b8 100644 --- a/pandas-stubs/core/indexers/objects.pyi +++ b/pandas-stubs/core/indexers/objects.pyi @@ -1,13 +1,20 @@ +from collections.abc import Sequence from typing import Any -import numpy as np -from pandas import DatetimeIndex +from pandas.core.indexes.datetimes import DatetimeIndex from pandas._libs.tslibs import BaseOffset +from pandas._typing import ( + np_1darray_intp, + np_ndarray_intp, +) class BaseIndexer: def __init__( - self, index_array: np.ndarray | None = None, window_size: int = 0, **kwargs: Any + self, + index_array: Sequence[float] | np_ndarray_intp | None = None, + window_size: int = 0, + **kwargs: Any, ) -> None: ... def get_window_bounds( self, @@ -16,14 +23,14 @@ class BaseIndexer: center: bool | None, closed: str | None = None, step: int | None = None, - ) -> tuple[np.ndarray, np.ndarray]: ... + ) -> tuple[np_1darray_intp, np_1darray_intp]: ... class FixedForwardWindowIndexer(BaseIndexer): ... class VariableOffsetWindowIndexer(BaseIndexer): def __init__( self, - index_array: np.ndarray | None = None, + index_array: np_ndarray_intp | None = None, window_size: int = 0, index: DatetimeIndex | None = None, offset: BaseOffset | None = None, diff --git a/pandas-stubs/core/indexers/utils.pyi b/pandas-stubs/core/indexers/utils.pyi index 99f1b018a..1913e15ee 100644 --- a/pandas-stubs/core/indexers/utils.pyi +++ b/pandas-stubs/core/indexers/utils.pyi @@ -3,7 +3,7 @@ from typing import overload from pandas._typing import ( AnyArrayLike, ListLike, - np_ndarray_int, + np_1darray_bool, ) @overload @@ -11,4 +11,4 @@ def check_array_indexer(array: AnyArrayLike, indexer: int) -> int: ... @overload def check_array_indexer(array: AnyArrayLike, indexer: slice) -> slice: ... @overload -def check_array_indexer(array: AnyArrayLike, indexer: ListLike) -> np_ndarray_int: ... +def check_array_indexer(array: AnyArrayLike, indexer: ListLike) -> np_1darray_bool: ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index b87872d78..69441edb2 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -46,6 +46,7 @@ from pandas._typing import ( TimeUnit, TimeZones, np_1darray, + np_1darray_bool, np_ndarray_bool, ) @@ -96,7 +97,7 @@ class _DatetimeFieldOps( ): ... _DTBoolOpsReturnType = TypeVar( - "_DTBoolOpsReturnType", bound=Series[bool] | np_1darray[np.bool] + "_DTBoolOpsReturnType", bound=Series[bool] | np_1darray_bool ) class _IsLeapYearProperty(Generic[_DTBoolOpsReturnType]): @@ -415,7 +416,7 @@ class DatetimeIndexProperties( Properties, _DatetimeNoTZProperties[ Index[int], - np_1darray[np.bool], + np_1darray_bool, DatetimeIndex, np_1darray[np.object_], np_1darray[np.object_], diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 8715182b5..1866e9464 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -107,6 +107,8 @@ from pandas._typing import ( TimedeltaDtypeArg, TimestampDtypeArg, np_1darray, + np_1darray_bool, + np_1darray_intp, np_ndarray, np_ndarray_anyint, np_ndarray_bool, @@ -340,7 +342,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): ) -> StringMethods[ # pyrefly: ignore[bad-specialization] Self, MultiIndex, - np_1darray[np.bool], + np_1darray_bool, Index[list[_str]], Index[int], Index[bytes], @@ -422,7 +424,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): def dropna(self, how: AnyAll = "any") -> Self: ... def unique(self, level: Hashable | None = None) -> Self: ... def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ... - def duplicated(self, keep: DropKeep = "first") -> np_1darray[np.bool]: ... + def duplicated(self, keep: DropKeep = "first") -> np_1darray_bool: ... def __and__(self, other: Never) -> Never: ... def __rand__(self, other: Never) -> Never: ... def __or__(self, other: Never) -> Never: ... @@ -446,14 +448,14 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): result_name: Hashable = ..., sort: bool | None = None, ) -> Self: ... - def get_loc(self, key: Label) -> int | slice | np_1darray[np.bool]: ... + def get_loc(self, key: Label) -> int | slice | np_1darray_bool: ... def get_indexer( self, target: Index, method: ReindexMethod | None = None, limit: int | None = None, tolerance: Scalar | AnyArrayLike | Sequence[Scalar] | None = None, - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... def reindex( self, target: Iterable[Any], @@ -461,7 +463,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): level: int | None = None, limit: int | None = None, tolerance: Scalar | AnyArrayLike | Sequence[Scalar] | None = None, - ) -> tuple[Index, np_1darray[np.intp] | None]: ... + ) -> tuple[Index, np_1darray_intp | None]: ... @overload def join( self, @@ -471,7 +473,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): level: Level | None = None, return_indexers: Literal[True], sort: bool = False, - ) -> tuple[Index, np_1darray[np.intp] | None, np_1darray[np.intp] | None]: ... + ) -> tuple[Index, np_1darray_intp | None, np_1darray_intp | None]: ... @overload def join( self, @@ -530,14 +532,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): ): ... @final def sort(self, *args: Any, **kwargs: Any) -> None: ... - def argsort(self, *args: Any, **kwargs: Any) -> np_1darray[np.intp]: ... + def argsort(self, *args: Any, **kwargs: Any) -> np_1darray_intp: ... def get_indexer_non_unique(self, target): ... @final def get_indexer_for(self, target, **kwargs: Any): ... - @final - def groupby(self, values) -> dict[Hashable, np.ndarray]: ... def map(self, mapper, na_action=...) -> Index: ... - def isin(self, values, level=...) -> np_1darray[np.bool]: ... + def isin(self, values, level=...) -> np_1darray_bool: ... def slice_indexer( self, start: Label | None = None, @@ -559,12 +559,12 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @property def shape(self) -> tuple[int, ...]: ... # Extra methods from old stubs - def __eq__(self, other: object) -> np_1darray[np.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] + def __eq__(self, other: object) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __ne__(self, other: object) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __le__(self, other: Self | S1) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __ge__(self, other: Self | S1) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __lt__(self, other: Self | S1) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __gt__(self, other: Self | S1) -> np_1darray_bool: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] @overload def __add__(self: Index[Never], other: _str) -> Never: ... @overload diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index 92c60cff9..930d78cbd 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/pandas-stubs/core/indexes/category.pyi @@ -3,8 +3,7 @@ from collections.abc import ( Iterable, ) -import numpy as np -from pandas.core import accessor +from pandas.core.accessor import PandasDelegate from pandas.core.arrays.categorical import Categorical from pandas.core.indexes.base import Index from pandas.core.indexes.extension import ExtensionIndex @@ -14,11 +13,12 @@ from pandas._typing import ( S1, Dtype, ListLike, + np_1darray_intp, ) -class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): - codes: np.ndarray = ... - categories: Index = ... +class CategoricalIndex(ExtensionIndex[S1], PandasDelegate): + codes: np_1darray_intp = ... + categories: Index[S1] = ... @property def array(self) -> Categorical: ... # type: ignore[override] # pyrefly: ignore[bad-override] def __new__( diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index 32739270e..ad220de4d 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -13,6 +13,7 @@ from pandas._typing import ( AxisIndex, GenericT_co, TimeUnit, + np_1darray_int64, np_ndarray_complex, ) @@ -62,3 +63,5 @@ class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1, GenericT_co]): @property def unit(self) -> TimeUnit: ... def as_unit(self, unit: TimeUnit) -> Self: ... + @property + def asi8(self) -> np_1darray_int64: ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index f0d1df255..662f4fc4f 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -34,7 +34,7 @@ from pandas._typing import ( IntervalClosedType, TimeUnit, TimeZones, - np_1darray, + np_1darray_intp, np_ndarray, np_ndarray_dt, np_ndarray_td, @@ -91,14 +91,14 @@ class DatetimeIndex( def inferred_type(self) -> str: ... def indexer_at_time( self, time: str | time, asof: bool = False - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... def indexer_between_time( self, start_time: time | str, end_time: time | str, include_start: bool = True, include_end: bool = True, - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... def to_julian_date(self) -> Index[float]: ... def isocalendar(self) -> DataFrame: ... @property diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 11f8227a1..193ceb08e 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -26,9 +26,12 @@ from pandas._typing import ( IntervalT, Label, MaskType, - np_1darray, + np_1darray_bool, np_ndarray_anyint, np_ndarray_bool, + np_ndarray_dt, + np_ndarray_int64, + np_ndarray_td, npt, ) @@ -36,7 +39,7 @@ from pandas.core.dtypes.dtypes import IntervalDtype as IntervalDtype _EdgesInt: TypeAlias = ( Sequence[int] - | npt.NDArray[np.int64] + | np_ndarray_int64 | npt.NDArray[np.int32] | npt.NDArray[np.intp] | pd.Series[int] @@ -46,16 +49,10 @@ _EdgesFloat: TypeAlias = ( Sequence[float] | npt.NDArray[np.float64] | pd.Series[float] | Index[float] ) _EdgesTimestamp: TypeAlias = ( - Sequence[DatetimeLike] - | npt.NDArray[np.datetime64] - | pd.Series[pd.Timestamp] - | pd.DatetimeIndex + Sequence[DatetimeLike] | np_ndarray_dt | pd.Series[pd.Timestamp] | pd.DatetimeIndex ) _EdgesTimedelta: TypeAlias = ( - Sequence[pd.Timedelta] - | npt.NDArray[np.timedelta64] - | pd.Series[pd.Timedelta] - | pd.TimedeltaIndex + Sequence[pd.Timedelta] | np_ndarray_td | pd.Series[pd.Timedelta] | pd.TimedeltaIndex ) _TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime _TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta @@ -216,7 +213,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): def memory_usage(self, deep: bool = False) -> int: ... @property def is_overlapping(self) -> bool: ... - def get_loc(self, key: Label) -> int | slice | np_1darray[np.bool]: ... + def get_loc(self, key: Label) -> int | slice | np_1darray_bool: ... def get_indexer_non_unique( self, target: Index ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ... @@ -248,7 +245,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): @overload # type: ignore[override] def __gt__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + ) -> np_1darray_bool: ... @overload def __gt__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -256,7 +253,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): @overload # type: ignore[override] def __ge__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + ) -> np_1darray_bool: ... @overload def __ge__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -264,7 +261,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): @overload # type: ignore[override] def __le__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + ) -> np_1darray_bool: ... @overload def __le__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: pd.Series[IntervalT] @@ -272,13 +269,13 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): @overload # type: ignore[override] def __lt__( self, other: IntervalT | IntervalIndex[IntervalT] - ) -> np_1darray[np.bool]: ... + ) -> np_1darray_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_1darray[np.bool]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + def __eq__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_1darray_bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] @overload def __eq__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[overload-overlap] @overload @@ -286,7 +283,7 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): self, other: object ) -> Literal[False]: ... @overload # type: ignore[override] - def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_1darray_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 7dc76e7f7..d44b42bd6 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -25,7 +25,7 @@ from pandas._typing import ( MaskType, NaPosition, SequenceNotStr, - np_1darray, + np_1darray_bool, np_ndarray_anyint, ) @@ -162,4 +162,4 @@ class MultiIndex(Index): def equal_levels(self, other): ... def insert(self, loc, item): ... def delete(self, loc): ... - def isin(self, values, level=...) -> np_1darray[np.bool]: ... + def isin(self, values, level=...) -> np_1darray_bool: ... diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 84abc62ea..96052bab8 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/pandas-stubs/core/indexes/period.pyi @@ -23,6 +23,8 @@ from pandas._typing import ( Dtype, Frequency, np_1darray, + np_1darray_intp, + np_ndarray_bool, ) class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexFieldOps): @@ -65,10 +67,8 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexField self, other: NaTType ) -> NaTType: ... def asof_locs( - self, - where: pd.DatetimeIndex | PeriodIndex, - mask: np_1darray[np.bool_], - ) -> np_1darray[np.intp]: ... + self, where: pd.DatetimeIndex | Self, mask: np_ndarray_bool + ) -> np_1darray_intp: ... @property def is_full(self) -> bool: ... @property diff --git a/pandas-stubs/core/indexes/range.pyi b/pandas-stubs/core/indexes/range.pyi index 8ffb2e0fa..b752fe2c5 100644 --- a/pandas-stubs/core/indexes/range.pyi +++ b/pandas-stubs/core/indexes/range.pyi @@ -22,7 +22,7 @@ from pandas._typing import ( HashableT, MaskType, Scalar, - np_1darray, + np_1darray_intp, np_ndarray_anyint, np_ndarray_bool, ) @@ -62,7 +62,7 @@ class RangeIndex(_IndexSubclassBase[int, np.int64]): def has_duplicates(self) -> bool: ... def factorize( self, sort: bool = False, use_na_sentinel: bool = True - ) -> tuple[np_1darray[np.intp], RangeIndex]: ... + ) -> tuple[np_1darray_intp, RangeIndex]: ... @property def size(self) -> int: ... def all(self, *args: Any, **kwargs: Any) -> bool: ... diff --git a/pandas-stubs/core/resample.pyi b/pandas-stubs/core/resample.pyi index 91299803e..54c748a67 100644 --- a/pandas-stubs/core/resample.pyi +++ b/pandas-stubs/core/resample.pyi @@ -12,16 +12,14 @@ from typing import ( ) import numpy as np -from pandas import ( - DataFrame, - Series, - Timedelta, -) +from pandas.core.frame import DataFrame from pandas.core.groupby.generic import SeriesGroupBy from pandas.core.groupby.groupby import BaseGroupBy from pandas.core.groupby.grouper import Grouper +from pandas.core.series import Series from typing_extensions import Self +from pandas._libs.tslibs.timedeltas import Timedelta from pandas._typing import ( S1, Axis, @@ -30,7 +28,7 @@ from pandas._typing import ( Scalar, TimeGrouperOrigin, TimestampConvention, - npt, + np_ndarray_float, ) _FrameGroupByFunc: TypeAlias = ( @@ -167,7 +165,7 @@ class Resampler(BaseGroupBy[NDFrameT]): @final def quantile( self, - q: float | list[float] | npt.NDArray[np.double] | Series[float] = 0.5, + q: float | list[float] | np_ndarray_float | Series[float] = 0.5, **kwargs: Any, ) -> NDFrameT: ... diff --git a/pandas-stubs/core/reshape/melt.pyi b/pandas-stubs/core/reshape/melt.pyi index 68a9ee57d..50c278e25 100644 --- a/pandas-stubs/core/reshape/melt.pyi +++ b/pandas-stubs/core/reshape/melt.pyi @@ -1,14 +1,19 @@ -from collections.abc import Hashable +from collections.abc import ( + Hashable, + Sequence, +) -import numpy as np from pandas.core.frame import DataFrame -from pandas._typing import HashableT +from pandas._typing import ( + HashableT, + np_ndarray, +) def melt( frame: DataFrame, - id_vars: tuple | list | np.ndarray | None = None, - value_vars: tuple | list | np.ndarray | None = None, + id_vars: Sequence[Hashable] | np_ndarray | None = None, + value_vars: Sequence[Hashable] | np_ndarray | None = None, var_name: str | None = None, value_name: Hashable = "value", col_level: int | str | None = None, diff --git a/pandas-stubs/core/reshape/pivot.pyi b/pandas-stubs/core/reshape/pivot.pyi index f328e1588..c7ea8fbd2 100644 --- a/pandas-stubs/core/reshape/pivot.pyi +++ b/pandas-stubs/core/reshape/pivot.pyi @@ -28,7 +28,7 @@ from pandas._typing import ( Label, Scalar, ScalarT, - npt, + np_ndarray, ) _PivotAggCallable: TypeAlias = Callable[[Series], ScalarT] @@ -86,7 +86,7 @@ def pivot_table( values: _PivotTableValuesTypes = None, *, index: Grouper, - columns: _PivotTableColumnsTypes | npt.NDArray[Any] | Index = None, + columns: _PivotTableColumnsTypes | np_ndarray | Index[Any] = None, aggfunc: ( _PivotAggFunc | Sequence[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc] ) = "mean", @@ -101,7 +101,7 @@ def pivot_table( def pivot_table( data: DataFrame, values: _PivotTableValuesTypes = None, - index: _PivotTableIndexTypes | npt.NDArray[Any] | Index = None, + index: _PivotTableIndexTypes | np_ndarray | Index[Any] = None, *, columns: Grouper, aggfunc: ( diff --git a/pandas-stubs/core/reshape/tile.pyi b/pandas-stubs/core/reshape/tile.pyi index 0eed14683..230bb50df 100644 --- a/pandas-stubs/core/reshape/tile.pyi +++ b/pandas-stubs/core/reshape/tile.pyi @@ -19,13 +19,16 @@ from pandas.core.series import Series from pandas._typing import ( IntervalT, Label, - npt, + np_1darray_float, + np_1darray_intp, + np_ndarray_anyint, + np_ndarray_float, ) @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + bins: int | Index[int] | Index[float] | Sequence[float] | Series, right: bool = ..., *, labels: Literal[False], @@ -34,10 +37,10 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[npt.NDArray[np.intp], npt.NDArray]: ... +) -> tuple[np_1darray_intp, np_1darray_float]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, bins: IntervalIndex[IntervalT], right: bool = ..., *, @@ -47,17 +50,11 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[npt.NDArray[np.intp], IntervalIndex[IntervalT]]: ... +) -> tuple[np_1darray_intp, IntervalIndex[IntervalT]]: ... @overload def cut( # pyright: ignore[reportOverlappingOverload] x: Series[Timestamp], - bins: ( - int - | Series[Timestamp] - | DatetimeIndex - | Sequence[Timestamp] - | Sequence[np.datetime64] - ), + bins: int | Series[Timestamp] | DatetimeIndex | Sequence[np.datetime64 | Timestamp], right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., *, @@ -83,7 +80,7 @@ def cut( @overload def cut( x: Series, - bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], + bins: int | Index[int] | Index[float] | Sequence[float] | Series, right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., *, @@ -92,7 +89,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[Series, npt.NDArray]: ... +) -> tuple[Series, np_1darray_float]: ... @overload def cut( x: Series, @@ -108,8 +105,8 @@ def cut( ) -> tuple[Series, IntervalIndex]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + bins: int | Index[int] | Index[float] | Sequence[float] | Series, right: bool = ..., labels: Sequence[Label] | None = ..., *, @@ -118,10 +115,10 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[Categorical, npt.NDArray]: ... +) -> tuple[Categorical, np_1darray_float]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, bins: IntervalIndex[IntervalT], right: bool = ..., labels: Sequence[Label] | None = ..., @@ -134,16 +131,8 @@ def cut( ) -> tuple[Categorical, IntervalIndex[IntervalT]]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: ( - int - | Series - | Index[int] - | Index[float] - | Sequence[int] - | Sequence[float] - | IntervalIndex - ), + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + bins: int | Sequence[float] | Index[int] | Index[float] | IntervalIndex | Series, right: bool = ..., *, labels: Literal[False], @@ -152,17 +141,16 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> npt.NDArray[np.intp]: ... +) -> np_1darray_intp: ... @overload def cut( x: Series[Timestamp], bins: ( int - | Series[Timestamp] - | DatetimeIndex - | Sequence[Timestamp] - | Sequence[np.datetime64] + | Sequence[np.datetime64 | Timestamp] | IntervalIndex[Interval[Timestamp]] + | DatetimeIndex + | Series[Timestamp] ), right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., @@ -175,15 +163,7 @@ def cut( @overload def cut( x: Series, - bins: ( - int - | Series - | Index[int] - | Index[float] - | Sequence[int] - | Sequence[float] - | IntervalIndex - ), + bins: int | Sequence[float] | Index[int] | Index[float] | IntervalIndex | Series, right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., retbins: Literal[False] = False, @@ -194,16 +174,8 @@ def cut( ) -> Series: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: ( - int - | Series - | Index[int] - | Index[float] - | Sequence[int] - | Sequence[float] - | IntervalIndex - ), + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + bins: int | Sequence[float] | Index[int] | Index[float] | IntervalIndex | Series, right: bool = ..., labels: Sequence[Label] | None = ..., retbins: Literal[False] = False, @@ -214,59 +186,70 @@ def cut( ) -> Categorical: ... @overload def qcut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, - *, + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], labels: Literal[False], retbins: Literal[False] = False, - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., -) -> npt.NDArray[np.intp]: ... + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", +) -> np_1darray_intp: ... @overload def qcut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, - labels: Sequence[Label] | None = ..., + x: Sequence[float] | np_ndarray_anyint | np_ndarray_float | Index, + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], + labels: Sequence[Label] | None = None, retbins: Literal[False] = False, - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", ) -> Categorical: ... @overload def qcut( x: Series, - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], labels: Literal[False] | Sequence[Label] | None = ..., retbins: Literal[False] = False, - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", ) -> Series: ... @overload def qcut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, - *, + x: ( + Sequence[float] + | np_ndarray_anyint + | np_ndarray_float + | Index[int] + | Index[float] + ), + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], labels: Literal[False], retbins: Literal[True], - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., -) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.double]]: ... + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", + # double when there are nan's +) -> tuple[np_1darray_intp | np_1darray_float, np_1darray_float]: ... @overload def qcut( - x: Series, - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, + x: Series[int] | Series[float], + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], labels: Literal[False] | Sequence[Label] | None = ..., *, retbins: Literal[True], - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., -) -> tuple[Series, npt.NDArray[np.double]]: ... + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", +) -> tuple[Series, np_1darray_float]: ... @overload def qcut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, + x: ( + Sequence[float] + | np_ndarray_anyint + | np_ndarray_float + | Index[int] + | Index[float] + ), + q: int | Sequence[float] | np_ndarray_float | Index[float] | Series[float], labels: Sequence[Label] | None = ..., *, retbins: Literal[True], - precision: int = ..., - duplicates: Literal["raise", "drop"] = ..., -) -> tuple[Categorical, npt.NDArray[np.double]]: ... + precision: int = 3, + duplicates: Literal["raise", "drop"] = "raise", +) -> tuple[Categorical, np_1darray_float]: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index f927835ce..61eb0222e 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -20,6 +20,7 @@ from datetime import ( timedelta, ) from pathlib import Path +import sys from typing import ( Any, ClassVar, @@ -47,6 +48,7 @@ from matplotlib.axes import ( SubplotBase, ) import numpy as np +from numpy import typing as npt from pandas import ( Index, Period, @@ -60,8 +62,12 @@ from pandas.core.api import ( Int32Dtype as Int32Dtype, Int64Dtype as Int64Dtype, ) +from pandas.core.arrays.base import ExtensionArray from pandas.core.arrays.boolean import BooleanDtype -from pandas.core.arrays.categorical import CategoricalAccessor +from pandas.core.arrays.categorical import ( + Categorical, + CategoricalAccessor, +) from pandas.core.arrays.datetimes import DatetimeArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.base import ( @@ -162,7 +168,6 @@ from pandas._typing import ( FloatFormatType, Frequency, GenericT, - GenericT_co, GroupByObjectNonScalar, HashableT1, IgnoreRaise, @@ -215,6 +220,11 @@ from pandas._typing import ( WriteBuffer, _T_co, np_1darray, + np_1darray_bool, + np_1darray_dt, + np_1darray_int64, + np_1darray_intp, + np_1darray_td, np_ndarray, np_ndarray_anyint, np_ndarray_bool, @@ -224,8 +234,6 @@ from pandas._typing import ( np_ndarray_num, np_ndarray_str, np_ndarray_td, - npt, - num, ) from pandas.core.dtypes.base import ExtensionDtype @@ -319,9 +327,9 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): value: S1 | ArrayLike | IndexOpsMixin[S1] | None, ) -> None: ... -_DataLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1] +_DataLike: TypeAlias = ArrayLike | dict[str, np_ndarray] | SequenceNotStr[S1] _DataLikeS1: TypeAlias = ( - ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | IndexOpsMixin[S1] + ArrayLike | dict[_str, np_ndarray] | Sequence[S1] | IndexOpsMixin[S1] ) class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @@ -533,8 +541,8 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @name.setter def name(self, value: Hashable | None) -> None: ... @property - def values(self) -> ArrayLike: ... - def ravel(self, order: _str = ...) -> np.ndarray: ... + def values(self) -> np_1darray | ExtensionArray | Categorical: ... + def ravel(self, order: _str = ...) -> np_1darray: ... def __len__(self) -> int: ... def view(self, dtype: Dtype | None = None) -> Series[S1]: ... @final @@ -874,13 +882,13 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def count(self) -> int: ... def mode(self, dropna: bool = True) -> Series[S1]: ... @overload - def unique(self: Series[Never]) -> np.ndarray: ... # type: ignore[overload-overlap] + def unique(self: Series[Never]) -> np_1darray: ... # type: ignore[overload-overlap] @overload def unique(self: Series[Timestamp]) -> DatetimeArray: ... # type: ignore[overload-overlap] @overload def unique(self: Series[Timedelta]) -> TimedeltaArray: ... # type: ignore[overload-overlap] @overload - def unique(self) -> np.ndarray: ... + def unique(self) -> np_1darray: ... @overload def drop_duplicates( self, @@ -957,27 +965,28 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def dot(self, other: DataFrame) -> Series[S1]: ... @overload def dot( - self, other: ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | Index[S1] - ) -> np.ndarray: ... + self, + other: ArrayLike | dict[_str, np_ndarray_num] | Sequence[S1] | Index[S1], + ) -> np_ndarray_num: ... @overload def __matmul__(self, other: Series) -> Scalar: ... @overload def __matmul__(self, other: DataFrame) -> Series: ... @overload - def __matmul__(self, other: np.ndarray) -> np.ndarray: ... + def __matmul__(self, other: np_ndarray_num) -> np_ndarray_num: ... @overload def __rmatmul__(self, other: Series) -> Scalar: ... @overload def __rmatmul__(self, other: DataFrame) -> Series: ... @overload - def __rmatmul__(self, other: np.ndarray) -> np.ndarray: ... + def __rmatmul__(self, other: np_ndarray_num) -> np_ndarray_num: ... @overload def searchsorted( self, value: ListLike, side: Literal["left", "right"] = ..., sorter: ListLike | None = None, - ) -> np_1darray[np.intp]: ... + ) -> np_1darray_intp: ... @overload def searchsorted( self, @@ -1395,7 +1404,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): n: int | None = None, frac: float | None = None, replace: _bool = False, - weights: _str | ListLike | np.ndarray | None = None, + weights: _str | ListLike | np_ndarray_float | None = None, random_state: RandomState | None = None, axis: AxisIndex | None = None, ignore_index: _bool = False, @@ -1610,7 +1619,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | np_ndarray_bool | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1626,7 +1635,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | np_ndarray_bool | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1642,7 +1651,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | np_ndarray_bool | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1658,7 +1667,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | np_ndarray_bool | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1670,12 +1679,13 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> Series[S1]: ... def case_when( self, - caselist: list[ + caselist: Sequence[ tuple[ Sequence[bool] + | np_1darray_bool | Series[bool] - | Callable[[Series], Series | np.ndarray | Sequence[bool]], - ListLikeU | Scalar | Callable[[Series], Series | np.ndarray], + | Callable[[Series], Sequence[bool] | np_1darray_bool | Series[bool]], + ListLikeU | Scalar | Callable[[Series], Series | np_ndarray], ], ], ) -> Series: ... @@ -2365,10 +2375,17 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): fill_value: float | None = None, axis: AxisIndex | None = 0, ) -> Series[int]: ... - @overload - def __rfloordiv__( # type: ignore[overload-overlap] - self: Series[Never], other: ScalarArrayIndexSeriesReal - ) -> Series: ... + if sys.version_info >= (3, 11): + @overload + def __rfloordiv__( # type: ignore[overload-overlap] + self: Series[Never], other: ScalarArrayIndexSeriesReal + ) -> Series: ... + else: + @overload + def __rfloordiv__( + self: Series[Never], other: ScalarArrayIndexSeriesReal + ) -> Series: ... + @overload def __rfloordiv__(self, other: np_ndarray_complex | np_ndarray_dt) -> Never: ... @overload @@ -2952,9 +2969,9 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): fill_value: float | None = None, axis: int = 0, ) -> Series[complex]: ... - def __mod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... + def __mod__(self, other: float | ListLike | Series[S1]) -> Series[S1]: ... def __ne__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __pow__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... + def __pow__(self, other: complex | ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] def __or__( # pyright: ignore[reportOverlappingOverload] # pyrefly: ignore[bad-override] @@ -2969,9 +2986,9 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> Series[bool]: ... @overload def __rand__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ... - def __rdivmod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __rmod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... - def __rpow__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... + def __rdivmod__(self, other: float | ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __rmod__(self, other: float | ListLike | Series[S1]) -> Series[S1]: ... + def __rpow__(self, other: complex | ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] def __ror__( # pyright: ignore[reportOverlappingOverload] # pyrefly: ignore[bad-override] @@ -4101,7 +4118,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> Series[S1]: ... def divmod( self, - other: num | ListLike | Series[S1], + other: float | ListLike | Series[S1], level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex = ..., @@ -4124,7 +4141,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): adjust: _bool = True, ignore_na: _bool = False, axis: Axis = 0, - times: np.ndarray | Series | None = None, + times: np_ndarray_dt | Series | None = None, method: CalculationMethod = "single", ) -> ExponentialMovingWindow[Series]: ... @final @@ -4259,7 +4276,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> S1: ... def mod( self, - other: num | ListLike | Series[S1], + other: float | ListLike | Series[S1], level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex | None = 0, @@ -4275,7 +4292,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def nunique(self, dropna: _bool = True) -> int: ... def pow( self, - other: num | ListLike | Series[S1], + other: complex | ListLike | Series[S1], level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex | None = 0, @@ -4420,7 +4437,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.datetime64]: ... + ) -> np_1darray_dt: ... @overload def to_numpy( self: Series[Timestamp], @@ -4436,7 +4453,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.timedelta64]: ... + ) -> np_1darray_td: ... @overload def to_numpy( self: Series[Timedelta], @@ -4460,7 +4477,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = False, na_value: Scalar = ..., **kwargs: Any, - ) -> np_1darray[np.int64]: ... + ) -> np_1darray_int64: ... @overload def to_numpy( self: Series[BaseOffset], @@ -4565,30 +4582,3 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> Self: ... @final def __bool__(self) -> NoReturn: ... - -@type_check_only -class _SeriesSubclassBase(Series[S1], Generic[S1, GenericT_co]): - @overload # type: ignore[override] - def to_numpy( # pyrefly: ignore - self, - dtype: None = None, - copy: bool = False, - na_value: Scalar = ..., - **kwargs: Any, - ) -> np_1darray[GenericT_co]: ... - @overload - def to_numpy( - self, - dtype: np.dtype[GenericT] | SupportsDType[GenericT] | type[GenericT], - copy: bool = False, - na_value: Scalar = ..., - **kwargs: Any, - ) -> np_1darray[GenericT]: ... - @overload - def to_numpy( # pyright: ignore[reportIncompatibleMethodOverride] - self, - dtype: DTypeLike, - copy: bool = False, - na_value: Scalar = ..., - **kwargs: Any, - ) -> np_1darray: ... diff --git a/pandas-stubs/core/strings/accessor.pyi b/pandas-stubs/core/strings/accessor.pyi index a290fa12f..24b3c2864 100644 --- a/pandas-stubs/core/strings/accessor.pyi +++ b/pandas-stubs/core/strings/accessor.pyi @@ -14,8 +14,6 @@ from typing import ( overload, ) -import numpy as np -import numpy.typing as npt import pandas as pd from pandas import ( DataFrame, @@ -31,7 +29,8 @@ from pandas._typing import ( DtypeObj, Scalar, T, - np_1darray, + np_1darray_bool, + np_ndarray_str, ) # Used for the result of str.split with expand=True @@ -39,7 +38,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_1darray[np.bool]) +_T_BOOL = TypeVar("_T_BOOL", bound=Series[bool] | np_1darray_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 @@ -67,9 +66,7 @@ class StringMethods( @overload def cat( self, - others: ( - Series[str] | Index[str] | pd.DataFrame | npt.NDArray[np.str_] | list[str] - ), + others: list[str] | np_ndarray_str | Series[str] | Index[str] | pd.DataFrame, sep: str | None = None, na_rep: str | None = None, join: AlignJoin = "left", diff --git a/pandas-stubs/core/tools/datetimes.pyi b/pandas-stubs/core/tools/datetimes.pyi index a4c841012..ebc25900a 100644 --- a/pandas-stubs/core/tools/datetimes.pyi +++ b/pandas-stubs/core/tools/datetimes.pyi @@ -26,7 +26,9 @@ from pandas._typing import ( IgnoreRaise, RaiseCoerce, TimestampConvertibleTypes, - npt, + np_ndarray_dt, + np_ndarray_int64, + np_ndarray_str, ) ArrayConvertible: TypeAlias = list | tuple | AnyArrayLike @@ -100,9 +102,9 @@ def to_datetime( Sequence[float | date] | list[str] | tuple[float | str | date, ...] - | npt.NDArray[np.datetime64] - | npt.NDArray[np.str_] - | npt.NDArray[np.int_] + | np_ndarray_dt + | np_ndarray_str + | np_ndarray_int64 | Index | ExtensionArray ), diff --git a/pandas-stubs/core/tools/numeric.pyi b/pandas-stubs/core/tools/numeric.pyi index 3e3d100ed..a90b07ddd 100644 --- a/pandas-stubs/core/tools/numeric.pyi +++ b/pandas-stubs/core/tools/numeric.pyi @@ -1,10 +1,11 @@ +from collections.abc import Sequence from typing import ( + Any, Literal, TypeAlias, overload, ) -import numpy as np import pandas as pd from pandas._libs.lib import _NoDefaultDoNotUse @@ -12,29 +13,39 @@ from pandas._typing import ( DtypeBackend, RaiseCoerce, Scalar, - npt, + SequenceNotStr, + np_1darray_anyint, + np_1darray_float, + np_ndarray, ) _Downcast: TypeAlias = Literal["integer", "signed", "unsigned", "float"] | None @overload -def to_numeric( +def to_numeric( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] arg: Scalar, - errors: Literal["raise", "coerce"] = ..., - downcast: _Downcast = ..., - dtype_backend: DtypeBackend | _NoDefaultDoNotUse = ..., + errors: RaiseCoerce = "raise", + downcast: _Downcast = None, + dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable", ) -> float: ... @overload def to_numeric( - arg: list | tuple | np.ndarray, - errors: RaiseCoerce = ..., - downcast: _Downcast = ..., - dtype_backend: DtypeBackend | _NoDefaultDoNotUse = ..., -) -> npt.NDArray: ... + arg: Sequence[int], + errors: RaiseCoerce = "raise", + downcast: _Downcast = None, + dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable", +) -> np_1darray_anyint: ... +@overload +def to_numeric( + arg: SequenceNotStr[Any] | np_ndarray, + errors: RaiseCoerce = "raise", + downcast: _Downcast = None, + dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable", +) -> np_1darray_float: ... @overload def to_numeric( arg: pd.Series, - errors: RaiseCoerce = ..., - downcast: _Downcast = ..., - dtype_backend: DtypeBackend | _NoDefaultDoNotUse = ..., + errors: RaiseCoerce = "raise", + downcast: _Downcast = None, + dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable", ) -> pd.Series: ... diff --git a/pandas-stubs/core/util/hashing.pyi b/pandas-stubs/core/util/hashing.pyi index fa59fc2ce..acb36d359 100644 --- a/pandas-stubs/core/util/hashing.pyi +++ b/pandas-stubs/core/util/hashing.pyi @@ -7,7 +7,7 @@ from pandas import ( from pandas._typing import ( ArrayLike, - npt, + np_1darray, ) def hash_pandas_object( @@ -16,10 +16,10 @@ def hash_pandas_object( encoding: str = "utf8", hash_key: str | None = "0123456789123456", categorize: bool = True, -) -> Series: ... +) -> Series[int]: ... def hash_array( vals: ArrayLike, encoding: str = "utf8", hash_key: str = "0123456789123456", categorize: bool = True, -) -> npt.NDArray[np.uint64]: ... +) -> np_1darray[np.uint64]: ... diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi index 8516c9075..29a461f8a 100644 --- a/pandas-stubs/io/formats/style.pyi +++ b/pandas-stubs/io/formats/style.pyi @@ -11,7 +11,6 @@ from typing import ( ) from matplotlib.colors import Colormap -import numpy as np from pandas.core.frame import DataFrame from pandas.core.series import Series @@ -31,7 +30,8 @@ from pandas._typing import ( T, WriteBuffer, WriteExcelBuffer, - npt, + np_ndarray, + np_ndarray_str, ) from pandas.io.excel import ExcelWriter @@ -52,7 +52,7 @@ class _SeriesFunc(Protocol): class _DataFrameFunc(Protocol): def __call__( self, series: DataFrame, /, *args: Any, **kwargs: Any - ) -> npt.NDArray | DataFrame: ... + ) -> np_ndarray | DataFrame: ... class _MapCallable(Protocol): def __call__( @@ -238,14 +238,14 @@ class Styler(StylerRenderer): @overload def apply( self, - func: _DataFrameFunc | Callable[[DataFrame], npt.NDArray | DataFrame], + func: _DataFrameFunc | Callable[[DataFrame], np_ndarray | DataFrame], axis: None, subset: Subset | None = ..., **kwargs: Any, ) -> Styler: ... def apply_index( self, - func: Callable[[Series], npt.NDArray[np.str_] | list[str] | Series[str]], + func: Callable[[Series], list[str] | np_ndarray_str | Series[str]], axis: Axis = ..., level: Level | list[Level] | None = ..., **kwargs: Any, @@ -295,7 +295,7 @@ class Styler(StylerRenderer): gmap: ( Sequence[float] | Sequence[Sequence[float]] - | npt.NDArray + | np_ndarray | DataFrame | Series | None @@ -313,7 +313,7 @@ class Styler(StylerRenderer): gmap: ( Sequence[float] | Sequence[Sequence[float]] - | npt.NDArray + | np_ndarray | DataFrame | Series | None @@ -334,7 +334,7 @@ class Styler(StylerRenderer): align: ( Literal["left", "right", "zero", "mid", "mean"] | float - | Callable[[Series | npt.NDArray | DataFrame], float] + | Callable[[Series | np_ndarray | DataFrame], float] ) = "mid", vmin: float | None = None, vmax: float | None = None, diff --git a/pandas-stubs/plotting/_core.pyi b/pandas-stubs/plotting/_core.pyi index f90e8f098..c5f4453a4 100644 --- a/pandas-stubs/plotting/_core.pyi +++ b/pandas-stubs/plotting/_core.pyi @@ -28,6 +28,7 @@ from pandas._typing import ( HashableT2, HashableT3, ListLikeHashable, + np_ndarray_float, npt, ) @@ -378,7 +379,7 @@ class PlotAccessor: | Callable[[gaussian_kde], float] | None ) = ..., - ind: npt.NDArray[np.double] | int | None = ..., + ind: np_ndarray_float | int | None = ..., *, subplots: Literal[False] | None = ..., **kwargs: Any, @@ -392,7 +393,7 @@ class PlotAccessor: | Callable[[gaussian_kde], float] | None ) = ..., - ind: npt.NDArray[np.double] | int | None = ..., + ind: np_ndarray_float | int | None = ..., *, subplots: Literal[True], **kwargs: Any, diff --git a/tests/__init__.py b/tests/__init__.py index 0e10b6b24..aa640305e 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -45,20 +45,45 @@ UIntDtypeArg as UIntDtypeArg, VoidDtypeArg as VoidDtypeArg, np_1darray as np_1darray, + np_1darray_anyint as np_1darray_anyint, + np_1darray_bool as np_1darray_bool, + np_1darray_dt as np_1darray_dt, + np_1darray_float as np_1darray_float, + np_1darray_int64 as np_1darray_int64, + np_1darray_intp as np_1darray_intp, + np_1darray_td as np_1darray_td, np_2darray as np_2darray, np_ndarray as np_ndarray, np_ndarray_bool as np_ndarray_bool, + np_ndarray_dt as np_ndarray_dt, np_ndarray_int as np_ndarray_int, + np_ndarray_intp as np_ndarray_intp, + np_ndarray_num as np_ndarray_num, + np_ndarray_str as np_ndarray_str, + np_ndarray_td as np_ndarray_td, ) else: _G = TypeVar("_G", bound=np.generic) _S = TypeVar("_S", bound=tuple[int, ...]) # Separately define here so pytest works np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]] + np_1darray_bool: TypeAlias = np_1darray[np.bool] + np_1darray_intp: TypeAlias = np_1darray[np.intp] + np_1darray_int64: TypeAlias = np_1darray[np.int64] + np_1darray_anyint: TypeAlias = np_1darray[np.integer] + np_1darray_float: TypeAlias = np_1darray[np.floating] + np_1darray_dt: TypeAlias = np_1darray[np.datetime64] + np_1darray_td: TypeAlias = np_1darray[np.timedelta64] 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_dt: TypeAlias = npt.NDArray[np.datetime64] np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger] + np_ndarray_intp: TypeAlias = npt.NDArray[np.intp] + np_ndarray_int64: TypeAlias = npt.NDArray[np.int64] + np_ndarray_num: TypeAlias = npt.NDArray[np.number] + np_ndarray_str: TypeAlias = npt.NDArray[np.str_] + np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64] TYPE_CHECKING_INVALID_USAGE: Final = TYPE_CHECKING WINDOWS = os.name == "nt" or "cygwin" in platform.system().lower() diff --git a/tests/extension/decimal/array.py b/tests/extension/decimal/array.py index 3a4d71047..b0ce317f6 100644 --- a/tests/extension/decimal/array.py +++ b/tests/extension/decimal/array.py @@ -11,7 +11,6 @@ import numbers import sys from typing import ( - TYPE_CHECKING, Any, cast, overload, @@ -53,8 +52,11 @@ pandas_dtype, ) -if TYPE_CHECKING: - from pandas._typing import np_1darray +from tests import ( + np_1darray, + np_1darray_bool, + np_ndarray, +) @register_extension_dtype @@ -94,7 +96,7 @@ class DecimalArray(OpsMixin, ExtensionScalarOpsMixin, ExtensionArray): def __init__( self, - values: MutableSequence[decimal._DecimalNew] | np.ndarray | ExtensionArray, + values: MutableSequence[decimal._DecimalNew] | np_ndarray | ExtensionArray, dtype: DecimalDtype | None = None, copy: bool = False, context: decimal.Context | None = None, @@ -126,7 +128,7 @@ def dtype(self) -> DecimalDtype: @classmethod def _from_sequence( cls, - scalars: list[decimal._DecimalNew] | np.ndarray | ExtensionArray, + scalars: list[decimal._DecimalNew] | np_ndarray | ExtensionArray, dtype: DecimalDtype | None = None, copy: bool = False, ) -> Self: @@ -144,7 +146,7 @@ def _from_sequence_of_strings( @classmethod def _from_factorized( cls, - values: list[decimal._DecimalNew] | np.ndarray | ExtensionArray, + values: list[decimal._DecimalNew] | np_ndarray | ExtensionArray, original: Any, ) -> Self: return cls(values) @@ -157,7 +159,7 @@ def to_numpy( copy: bool = False, na_value: object = no_default, decimals: int | None = None, - ) -> np.ndarray: + ) -> np_ndarray: result = np.asarray(self, dtype=dtype) if decimals is not None: result = np.asarray([round(x, decimals) for x in result]) @@ -192,7 +194,7 @@ def reconstruct( decimal.Decimal | numbers.Number | list[decimal._DecimalNew] - | np.ndarray + | np_ndarray ), ) -> decimal.Decimal | numbers.Number | DecimalArray: if isinstance(x, (decimal.Decimal, numbers.Number)): @@ -281,7 +283,7 @@ def nbytes(self) -> int: return n * sys.getsizeof(self[0]) return 0 - def isna(self) -> np_1darray[np.bool_]: + def isna(self) -> np_1darray_bool: if sys.version_info < (3, 11): return np.array([x.is_nan() for x in self._data], bool) # type: ignore[return-value] # pyright: ignore[reportReturnType] return np.array([x.is_nan() for x in self._data], bool) @@ -320,7 +322,7 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs: Any) -> Any: def _cmp_method( self, other: Any, op: Callable[[Self, ExtensionArray | list[Any]], bool] - ) -> np_1darray[np.bool_]: + ) -> np_1darray_bool: # For use with OpsMixin def convert_values(param: Any) -> ExtensionArray | list[Any]: if isinstance(param, ExtensionArray) or is_list_like(param): @@ -337,9 +339,7 @@ def convert_values(param: Any) -> ExtensionArray | list[Any]: # a TypeError should be raised res = [op(a, b) for (a, b) in zip(lvalues, rvalues)] - return cast( - np.ndarray[tuple[int], np.dtype[np.bool_]], np.asarray(res, dtype=bool) - ) + return cast(np_1darray_bool, np.asarray(res, dtype=bool)) def value_counts(self, dropna: bool = True) -> Series: from pandas.core.algorithms import value_counts diff --git a/tests/indexes/test_datetime_index.py b/tests/indexes/test_datetime_index.py index c296efd9a..ec6932028 100644 --- a/tests/indexes/test_datetime_index.py +++ b/tests/indexes/test_datetime_index.py @@ -11,6 +11,7 @@ from tests import ( check, np_1darray, + np_1darray_int64, ) @@ -130,3 +131,8 @@ def test_datetimeindex_snap() -> None: ), pd.DatetimeIndex, ) + + +def test_datetimeindex_properties() -> None: + dti = pd.date_range("2023-01-01", "2023-02-01") + check(assert_type(dti.asi8, np_1darray_int64), np_1darray_int64, np.integer) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 92df8721f..8903c5666 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -9,7 +9,6 @@ ) import numpy as np -from numpy import typing as npt import pandas as pd from pandas.core.arrays.base import ExtensionArray from pandas.core.arrays.categorical import Categorical @@ -29,6 +28,7 @@ TYPE_CHECKING_INVALID_USAGE, check, np_1darray, + np_ndarray_dt, pytest_warns_bounded, ) @@ -629,7 +629,7 @@ def test_interval_index_arrays() -> None: pd.IntervalIndex, pd.Interval, ) - left_dt64_arr: npt.NDArray[np.datetime64] = np.array( + left_dt64_arr: np_ndarray_dt = np.array( [ np.datetime64("2000-01-01"), np.datetime64("2001-01-01"), @@ -638,7 +638,7 @@ def test_interval_index_arrays() -> None: ], dtype="datetime64[ns]", ) - right_dt_arr: npt.NDArray[np.datetime64] = np.array( + right_dt_arr: np_ndarray_dt = np.array( [ np.datetime64("2001-01-01"), np.datetime64("2002-01-01"), diff --git a/tests/indexes/test_timedelta_indexes.py b/tests/indexes/test_timedelta_indexes.py new file mode 100644 index 000000000..daa7e8cf6 --- /dev/null +++ b/tests/indexes/test_timedelta_indexes.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +from typing_extensions import ( + assert_type, +) + +from tests import ( + check, + np_1darray_int64, +) + + +@pytest.fixture +def tdi() -> pd.TimedeltaIndex: + idx = pd.timedelta_range("1 days", periods=3, freq="D") + return check(assert_type(idx, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + + +def test_timedelta_index_properties(tdi: pd.TimedeltaIndex) -> None: + check(assert_type(tdi.asi8, np_1darray_int64), np_1darray_int64, np.integer) diff --git a/tests/scalars/test_scalars.py b/tests/scalars/test_scalars.py index 61e3fd684..fd954051a 100644 --- a/tests/scalars/test_scalars.py +++ b/tests/scalars/test_scalars.py @@ -27,7 +27,10 @@ check, np_1darray, np_2darray, + np_ndarray, np_ndarray_bool, + np_ndarray_dt, + np_ndarray_td, pytest_warns_bounded, ) @@ -504,12 +507,8 @@ def test_timedelta_properties_methods() -> None: def test_timedelta_add_sub() -> None: td = pd.Timedelta("1 day") - ndarray_td64: npt.NDArray[np.timedelta64] = np.array( - [1, 2, 3], dtype="timedelta64[D]" - ) - ndarray_dt64: npt.NDArray[np.datetime64] = np.array( - [1, 2, 3], dtype="datetime64[D]" - ) + ndarray_td64: np_ndarray_td = np.array([1, 2, 3], dtype="timedelta64[D]") + ndarray_dt64: np_ndarray_dt = np.array([1, 2, 3], dtype="datetime64[D]") as_period = pd.Period("2012-01-01", freq="D") as_timestamp = pd.Timestamp("2012-01-01") as_datetime = dt.datetime(2012, 1, 1) @@ -535,16 +534,8 @@ def test_timedelta_add_sub() -> None: check(assert_type(td + as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td + as_period_index, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(td + as_datetime_index, pd.DatetimeIndex), pd.DatetimeIndex) - check( - assert_type(td + as_ndarray_td64, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) - check( - assert_type(td + as_ndarray_dt64, npt.NDArray[np.datetime64]), - np.ndarray, - np.datetime64, - ) + check(assert_type(td + as_ndarray_td64, np_ndarray_td), np_ndarray, np.timedelta64) + check(assert_type(td + as_ndarray_dt64, np_ndarray_dt), np_ndarray, np.datetime64) check(assert_type(td + as_nat, NaTType), NaTType) check(assert_type(as_period + td, pd.Period), pd.Period) @@ -572,14 +563,7 @@ def test_timedelta_add_sub() -> None: check(assert_type(as_timedelta_index + td, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(as_period_index + td, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(as_datetime_index + td, pd.DatetimeIndex), pd.DatetimeIndex) - check( - assert_type( - as_ndarray_td64 + td, - npt.NDArray[np.timedelta64], - ), - np.ndarray, - np.timedelta64, - ) + check(assert_type(as_ndarray_td64 + td, np_ndarray_td), np_ndarray, np.timedelta64) check(assert_type(as_nat + td, NaTType), NaTType) # sub is not symmetric with dates. In general date_like - timedelta is @@ -600,11 +584,7 @@ def test_timedelta_add_sub() -> None: check(assert_type(td - as_dt_timedelta, pd.Timedelta), pd.Timedelta) check(assert_type(td - as_timedelta64, pd.Timedelta), pd.Timedelta) check(assert_type(td - as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) - check( - assert_type(td - as_ndarray_td64, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) + check(assert_type(td - as_ndarray_td64, np_ndarray_td), np_ndarray, np.timedelta64) check(assert_type(td - as_nat, NaTType), NaTType) check(assert_type(as_period - td, pd.Period), pd.Period) check(assert_type(as_timestamp - td, pd.Timestamp), pd.Timestamp) @@ -631,22 +611,8 @@ def test_timedelta_add_sub() -> None: check(assert_type(as_timedelta_index - td, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(as_period_index - td, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(as_datetime_index - td, pd.DatetimeIndex), pd.DatetimeIndex) - check( - assert_type( - as_ndarray_td64 - td, - npt.NDArray[np.timedelta64], - ), - np.ndarray, - np.timedelta64, - ) - check( - assert_type( - as_ndarray_dt64 - td, - npt.NDArray[np.datetime64], - ), - np.ndarray, - np.datetime64, - ) + check(assert_type(as_ndarray_td64 - td, np_ndarray_td), np_ndarray, np.timedelta64) + check(assert_type(as_ndarray_dt64 - td, np_ndarray_dt), np_ndarray, np.datetime64) check(assert_type(as_nat - td, NaTType), NaTType) @@ -663,43 +629,21 @@ def test_timedelta_mul_div() -> None: check(assert_type(td * md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td * md_float, pd.Timedelta), pd.Timedelta) - check( - assert_type(td * md_ndarray_intp, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) - check( - assert_type(td * md_ndarray_float, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) + check(assert_type(td * md_ndarray_intp, np_ndarray_td), np_ndarray, np.timedelta64) + check(assert_type(td * md_ndarray_float, np_ndarray_td), np_ndarray, np.timedelta64) check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta) check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta) - check( - assert_type(md_ndarray_intp * td, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) - check( - assert_type(md_ndarray_float * td, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) + check(assert_type(md_ndarray_intp * td, np_ndarray_td), np_ndarray, np.timedelta64) + check(assert_type(md_ndarray_float * td, np_ndarray_td), np_ndarray, np.timedelta64) check(assert_type(td // td, int), int) check(assert_type(td // pd.NaT, float), float) check(assert_type(td // md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td // md_float, pd.Timedelta), pd.Timedelta) + check(assert_type(td // md_ndarray_intp, np_ndarray_td), np_ndarray, np.timedelta64) check( - assert_type(td // md_ndarray_intp, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) - check( - assert_type(td // md_ndarray_float, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, + assert_type(td // md_ndarray_float, np_ndarray_td), np_ndarray, np.timedelta64 ) check(assert_type(pd.NaT // td, float), float) @@ -716,16 +660,8 @@ def test_timedelta_mul_div() -> None: check(assert_type(td / pd.NaT, float), float) check(assert_type(td / md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td / md_float, pd.Timedelta), pd.Timedelta) - check( - assert_type(td / md_ndarray_intp, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) - check( - assert_type(td / md_ndarray_float, npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, - ) + check(assert_type(td / md_ndarray_intp, np_ndarray_td), np_ndarray, np.timedelta64) + check(assert_type(td / md_ndarray_float, np_ndarray_td), np_ndarray, np.timedelta64) check(assert_type(pd.NaT / td, float), float) # Note: None of the reverse truediv work @@ -748,13 +684,11 @@ def test_timedelta_mod_abs_unary() -> None: check(assert_type(td % 3.5, pd.Timedelta), pd.Timedelta) check(assert_type(td % td, pd.Timedelta), pd.Timedelta) check( - assert_type(td % np.array([1, 2, 3]), npt.NDArray[np.timedelta64]), - np.ndarray, - np.timedelta64, + assert_type(td % np.array([1, 2, 3]), np_ndarray_td), np_ndarray, np.timedelta64 ) check( - assert_type(td % np.array([1.2, 2.2, 3.4]), npt.NDArray[np.timedelta64]), - np.ndarray, + assert_type(td % np.array([1.2, 2.2, 3.4]), np_ndarray_td), + np_ndarray, np.timedelta64, ) int_series = pd.Series([1, 2, 3], dtype=int) @@ -916,11 +850,11 @@ def test_timedelta_cmp_array() -> None: 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) + gt_nd1 = check(assert_type(td > arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd1 = check(assert_type(td <= arr_nd, np_ndarray_bool), np_ndarray_bool, 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) + gt_nd2 = check(assert_type(arr_nd > td, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd2 = check(assert_type(arr_nd <= td, np_ndarray_bool), np_ndarray_bool, 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]) @@ -936,11 +870,11 @@ def test_timedelta_cmp_array() -> None: 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) + lt_nd1 = check(assert_type(td < arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd1 = check(assert_type(td >= arr_nd, np_ndarray_bool), np_ndarray_bool, 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) + lt_nd2 = check(assert_type(arr_nd < td, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd2 = check(assert_type(arr_nd >= td, np_ndarray_bool), np_ndarray_bool, 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]) @@ -956,8 +890,8 @@ def test_timedelta_cmp_array() -> None: 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) + eq_nd1 = check(assert_type(td == arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ne_nd1 = check(assert_type(td != arr_nd, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1087,7 +1021,7 @@ def test_timestamp_properties() -> None: def test_timestamp_add_sub() -> None: ts = pd.Timestamp("2000-1-1") - np_td64_arr: npt.NDArray[np.timedelta64] = np.array([1, 2], dtype="timedelta64[ns]") + np_td64_arr: np_ndarray_td = np.array([1, 2], dtype="timedelta64[ns]") as_pd_timedelta = pd.Timedelta(days=1) as_dt_timedelta = dt.timedelta(days=1) @@ -1126,17 +1060,10 @@ def test_timestamp_add_sub() -> None: ) check( - assert_type(ts + as_np_ndarray_td64, npt.NDArray[np.datetime64]), - np.ndarray, - np.datetime64, + assert_type(ts + as_np_ndarray_td64, np_ndarray_dt), np_ndarray, np.datetime64 ) check( - assert_type( - as_np_ndarray_td64 + ts, - npt.NDArray[np.datetime64], - ), - np.ndarray, - np.datetime64, + assert_type(as_np_ndarray_td64 + ts, np_ndarray_dt), np_ndarray, np.datetime64 ) # Reverse order is not possible for all of these @@ -1150,9 +1077,7 @@ def test_timestamp_add_sub() -> None: pd.Timestamp, ) check( - assert_type(ts - as_np_ndarray_td64, npt.NDArray[np.datetime64]), - np.ndarray, - np.datetime64, + assert_type(ts - as_np_ndarray_td64, np_ndarray_dt), np_ndarray, np.datetime64 ) @@ -1321,11 +1246,11 @@ def test_timestamp_cmp_array() -> None: 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) + gt_nd1 = check(assert_type(ts > arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd1 = check(assert_type(ts <= arr_nd, np_ndarray_bool), np_ndarray_bool, 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) + gt_nd2 = check(assert_type(arr_nd > ts, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd2 = check(assert_type(arr_nd <= ts, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1341,11 +1266,11 @@ def test_timestamp_cmp_array() -> None: 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) + lt_nd1 = check(assert_type(ts < arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd1 = check(assert_type(ts >= arr_nd, np_ndarray_bool), np_ndarray_bool, 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) + lt_nd2 = check(assert_type(arr_nd < ts, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd2 = check(assert_type(arr_nd >= ts, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1361,8 +1286,8 @@ def test_timestamp_cmp_array() -> None: 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) + eq_nd1 = check(assert_type(ts == arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ne_nd1 = check(assert_type(ts != arr_nd, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1925,8 +1850,8 @@ def test_period_cmp_array() -> None: 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) + gt_nd1 = check(assert_type(p > arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd1 = check(assert_type(p <= arr_nd, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1935,8 +1860,8 @@ def test_period_cmp_array() -> None: 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) + gt_nd2 = check(assert_type(arr_nd > p, np_ndarray_bool), np_ndarray_bool, np.bool) + le_nd2 = check(assert_type(arr_nd <= p, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1946,8 +1871,8 @@ def test_period_cmp_array() -> None: 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) + lt_nd1 = check(assert_type(p < arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd1 = check(assert_type(p >= arr_nd, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1956,8 +1881,8 @@ def test_period_cmp_array() -> None: 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) + lt_nd2 = check(assert_type(arr_nd < p, np_ndarray_bool), np_ndarray_bool, np.bool) + ge_nd2 = check(assert_type(arr_nd >= p, np_ndarray_bool), np_ndarray_bool, 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]) @@ -1967,8 +1892,8 @@ def test_period_cmp_array() -> None: 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) + eq_nd1 = check(assert_type(p == arr_nd, np_ndarray_bool), np_ndarray_bool, np.bool) + ne_nd1 = check(assert_type(p != arr_nd, np_ndarray_bool), np_ndarray_bool, 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]) diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 474649999..40c3ec10a 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import ( + Callable, Hashable, Iterable, Iterator, @@ -33,6 +34,7 @@ ) from pandas.api.typing import NAType from pandas.core.arrays.datetimes import DatetimeArray +from pandas.core.arrays.string_ import StringArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.window import ExponentialMovingWindow from pandas.core.window.expanding import Expanding @@ -59,6 +61,8 @@ check, ensure_clean, np_1darray, + np_1darray_bool, + np_ndarray_num, pytest_warns_bounded, ) from tests.extension.decimal.array import DecimalDtype @@ -771,7 +775,7 @@ def test_types_value_counts() -> None: def test_types_unique() -> None: s = pd.Series([-10, 2, 2, 3, 10, 10]) - check(assert_type(s.unique(), np.ndarray), np.ndarray) + check(assert_type(s.unique(), np_1darray), np_1darray) def test_types_apply() -> None: @@ -1410,31 +1414,41 @@ def test_types_rename_axis() -> None: def test_types_values() -> None: check( - assert_type(pd.Series([1, 2, 3]).values, ExtensionArray | np.ndarray), - np.ndarray, + assert_type( + pd.Series([1, 2, 3]).values, + np_1darray | ExtensionArray | pd.Categorical, + ), + np_1darray, + np.integer, ) - valresult_type: type[np.ndarray | ExtensionArray] + valresult_type: type[np_1darray | ExtensionArray | pd.Categorical] if PD_LTE_23: - valresult_type = np.ndarray + valresult_type = np_1darray else: - valresult_type = ExtensionArray + valresult_type = StringArray check( - assert_type(pd.Series(list("aabc")).values, np.ndarray | ExtensionArray), + assert_type( + pd.Series(list("aabc")).values, + np_1darray | ExtensionArray | pd.Categorical, + ), valresult_type, + str, ) check( assert_type( pd.Series(list("aabc")).astype("category").values, - np.ndarray | ExtensionArray, + np_1darray | ExtensionArray | pd.Categorical, ), pd.Categorical, + str, ) check( assert_type( pd.Series(pd.date_range("20130101", periods=3, tz="US/Eastern")).values, - np.ndarray | ExtensionArray, + np_1darray | ExtensionArray | pd.Categorical, ), - np.ndarray, + np_1darray, + np.datetime64, ) @@ -1570,8 +1584,8 @@ def test_types_dot() -> None: check(assert_type(s1 @ s2, Scalar), np.int64) check(assert_type(s1.dot(df1), "pd.Series[int]"), pd.Series, np.int64) check(assert_type(s1 @ df1, pd.Series), pd.Series) - check(assert_type(s1.dot(n1), np.ndarray), np.ndarray) - check(assert_type(s1 @ n1, np.ndarray), np.ndarray) + check(assert_type(s1.dot(n1), np_ndarray_num), np.ndarray) + check(assert_type(s1 @ n1, np_ndarray_num), np.ndarray) def test_series_loc_setitem() -> None: @@ -3611,13 +3625,20 @@ def test_case_when() -> None: c = pd.Series([6, 7, 8, 9], name="c") a = pd.Series([0, 0, 1, 2]) b = pd.Series([0, 3, 4, 5]) - r = c.case_when( - caselist=[ - (a.gt(0), a), - (b.gt(0), b), - ] - ) - check(assert_type(r, pd.Series), pd.Series) + + c0 = [(a.gt(0), a), (b.gt(0), b)] + check(assert_type(c.case_when(c0), pd.Series), pd.Series) + + def foo_factory( + thresh: int, + ) -> Callable[[pd.Series], pd.Series[bool] | np_1darray_bool]: + def foo(s: pd.Series) -> pd.Series[bool] | np_1darray_bool: + return s >= thresh + + return foo + + c1 = [(foo_factory(2), a), (foo_factory(0), b)] + check(assert_type(c.case_when(c1), pd.Series), pd.Series) def test_series_unique_timestamp() -> None: diff --git a/tests/test_extension.py b/tests/test_extension.py index a12f264d2..0f8653386 100644 --- a/tests/test_extension.py +++ b/tests/test_extension.py @@ -3,9 +3,13 @@ import numpy as np import pandas as pd from pandas.arrays import IntegerArray +from pandas.core.indexers import check_array_indexer from typing_extensions import assert_type -from tests import check +from tests import ( + check, + np_1darray_bool, +) from tests.extension.decimal.array import ( DecimalArray, DecimalDtype, @@ -38,3 +42,17 @@ def test_ExtensionArray_reduce_accumulate() -> None: ) check(assert_type(_data._reduce("max"), object), np.integer) check(assert_type(_data._accumulate("cumsum"), IntegerArray), IntegerArray) + + +def test_array_indexer() -> None: + arr = pd.array([1, 2]) + + m_pd = pd.array([True, False]) + check(assert_type(check_array_indexer(arr, m_pd), np_1darray_bool), np_1darray_bool) + + m_np = np.array([True, False], np.bool_) + check(assert_type(check_array_indexer(arr, m_np), np_1darray_bool), np_1darray_bool) + + check(assert_type(check_array_indexer(arr, 1), int), int) + + check(assert_type(check_array_indexer(arr, slice(0, 1, 1)), slice), slice) diff --git a/tests/test_frame.py b/tests/test_frame.py index 6a3973905..eb2960b0f 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -54,7 +54,9 @@ TYPE_CHECKING_INVALID_USAGE, check, ensure_clean, + np_1darray, np_2darray, + np_ndarray, pytest_warns_bounded, ) @@ -630,7 +632,7 @@ def test_types_eval() -> None: check(assert_type(df.eval("C = col1 % col2 == 0", inplace=True), None), type(None)) check( assert_type( - df.eval("E = col1 > col2"), Scalar | np.ndarray | pd.DataFrame | pd.Series + df.eval("E = col1 > col2"), Scalar | np_ndarray | pd.DataFrame | pd.Series ), pd.DataFrame, ) @@ -1086,7 +1088,7 @@ def test_types_value_counts() -> None: def test_types_unique() -> None: # This is really more for of a Series test df = pd.DataFrame(data={"col1": [1, 2], "col2": [1, 4]}) - check(assert_type(df["col1"].unique(), np.ndarray), np.ndarray) + check(assert_type(df["col1"].unique(), np_1darray), np_1darray) def test_types_apply() -> None: diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 6925582ac..856084fe2 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -9,14 +9,14 @@ ) import numpy as np -import numpy.typing as npt import pandas as pd -from pandas import Grouper from pandas.api.extensions import ExtensionArray from pandas.api.typing import ( NaTType, NAType, ) +from pandas.core.arrays.integer import IntegerArray +from pandas.core.groupby.grouper import Grouper import pandas.util as pdutil # TODO: pandas-dev/pandas#55023 @@ -33,7 +33,14 @@ TYPE_CHECKING_INVALID_USAGE, check, np_1darray, + np_1darray_anyint, + np_1darray_dt, + np_1darray_float, + np_1darray_int64, + np_1darray_intp, + np_1darray_td, np_2darray, + np_ndarray, np_ndarray_bool, pytest_warns_bounded, ) @@ -72,23 +79,16 @@ def test_types_to_datetime() -> None: pd.DatetimeIndex, ) check( - assert_type( - pd.to_datetime([1, 2], unit="D", origin=3), - pd.DatetimeIndex, - ), + assert_type(pd.to_datetime([1, 2], unit="D", origin=3), pd.DatetimeIndex), pd.DatetimeIndex, ) check( - assert_type( - pd.to_datetime(["2022-01-03", "2022-02-22"]), - pd.DatetimeIndex, - ), + assert_type(pd.to_datetime(["2022-01-03", "2022-02-22"]), pd.DatetimeIndex), pd.DatetimeIndex, ) check( assert_type( - pd.to_datetime(pd.Index(["2022-01-03", "2022-02-22"])), - pd.DatetimeIndex, + pd.to_datetime(pd.Index(["2022-01-03", "2022-02-22"])), pd.DatetimeIndex ), pd.DatetimeIndex, ) @@ -131,9 +131,10 @@ def test_types_concat_none() -> None: pd.DataFrame, ) - if TYPE_CHECKING_INVALID_USAGE: - # using assert_type as otherwise the second call would not be type-checked + def _0() -> None: # pyright: ignore[reportUnusedFunction] assert_type(pd.concat({"a": None}), Never) + + def _1() -> None: # pyright: ignore[reportUnusedFunction] assert_type(pd.concat([None]), Never) @@ -196,13 +197,7 @@ def test_types_concat() -> None: check(assert_type(ts1, "pd.Series[int]"), pd.Series, np.integer) check(assert_type(ts2, "pd.Series[int]"), pd.Series, np.integer) - check( - assert_type( - pd.concat({1: s, None: s2}, axis=1), - pd.DataFrame, - ), - pd.DataFrame, - ) + check(assert_type(pd.concat({1: s, None: s2}, axis=1), pd.DataFrame), pd.DataFrame) df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) df2 = pd.DataFrame(data={"col1": [10, 20], "col2": [30, 40]}) @@ -567,66 +562,55 @@ def test_read_xml() -> None: def test_unique() -> None: # Taken from the docs + ints = [2, 1, 3, 3] check( - assert_type(pd.unique(pd.Series([2, 1, 3, 3])), np.ndarray | ExtensionArray), - np.ndarray, + assert_type(pd.unique(pd.Series(ints)), np_1darray | ExtensionArray), + np_1darray_int64, ) - check( - assert_type(pd.unique(pd.Series([2] + [1] * 5)), np.ndarray | ExtensionArray), - np.ndarray, + assert_type( + pd.unique(pd.Series(ints, dtype="Int64")), np_1darray | ExtensionArray + ), + IntegerArray, + np.integer, ) check( - assert_type( - pd.unique(pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")])), - np.ndarray | ExtensionArray, - ), - np.ndarray, + assert_type(pd.unique(pd.Series([2] + [1] * 5)), np_1darray | ExtensionArray), + np_1darray_int64, ) + unzoned_timestamps = [pd.Timestamp("20160101"), pd.Timestamp("20160101")] check( assert_type( - pd.unique( - pd.Series( - [ - pd.Timestamp("20160101", tz="US/Eastern"), - pd.Timestamp("20160101", tz="US/Eastern"), - ] - ) - ), - np.ndarray | ExtensionArray, + pd.unique(pd.Series(unzoned_timestamps)), np_1darray | ExtensionArray ), - pd.arrays.DatetimeArray, + np_1darray, + np.datetime64, ) + + zoned_timestamps = [ts.tz_localize("US/Eastern") for ts in unzoned_timestamps] check( assert_type( - pd.unique( - pd.Index( - [ - pd.Timestamp("20160101", tz="US/Eastern"), - pd.Timestamp("20160101", tz="US/Eastern"), - ] - ) - ), - np.ndarray, + pd.unique(pd.Series(zoned_timestamps)), np_1darray | ExtensionArray ), - pd.DatetimeIndex, + pd.arrays.DatetimeArray, + pd.Timestamp, ) - check(assert_type(pd.unique(np.array(list("baabc"))), np.ndarray), np.ndarray) + check(assert_type(pd.unique(np.array(list("baabc"))), np_1darray), np_1darray) check( assert_type( pd.unique(pd.Series(pd.Categorical(list("baabc")))), - np.ndarray | ExtensionArray, + np_1darray | ExtensionArray, ), pd.Categorical, ) check( assert_type( pd.unique(pd.Series(pd.Categorical(list("baabc"), categories=list("abc")))), - np.ndarray | ExtensionArray, + np_1darray | ExtensionArray, ), pd.Categorical, ) @@ -637,36 +621,10 @@ def test_unique() -> None: pd.Categorical(list("baabc"), categories=list("abc"), ordered=True) ) ), - np.ndarray | ExtensionArray, + np_1darray | ExtensionArray, ), pd.Categorical, ) - check( - assert_type(pd.unique(pd.Index(["a", "b", "c", "a"])), np.ndarray), - np.ndarray if PD_LTE_23 else pd.Index, - ) - check( - assert_type(pd.unique(pd.RangeIndex(0, 10)), np.ndarray), - np.ndarray if PD_LTE_23 else pd.Index, - ) - check( - assert_type(pd.unique(pd.Categorical(["a", "b", "c", "a"])), pd.Categorical), - pd.Categorical, - ) - check( - assert_type( - pd.unique(pd.period_range("2001Q1", periods=10, freq="D")), - pd.PeriodIndex, - ), - pd.PeriodIndex, - ) - check( - assert_type( - pd.unique(pd.timedelta_range(start="1 day", periods=4)), - np.ndarray, - ), - np.ndarray if PD_LTE_23 else pd.Index, - ) # GH 200 @@ -702,29 +660,44 @@ def test_arrow_dtype() -> None: def test_hashing() -> None: a = np.array([1, 2, 3]) - check(assert_type(pdutil.hash_array(a), npt.NDArray[np.uint64]), np.ndarray) + check( + assert_type(pdutil.hash_array(a), np_1darray[np.uint64]), np_1darray[np.uint64] + ) check( assert_type( pdutil.hash_array(a, encoding="latin1", hash_key="1", categorize=True), - npt.NDArray[np.uint64], + np_1darray[np.uint64], ), - np.ndarray, + np_1darray[np.uint64], ) b = pd.Series(a) c = pd.DataFrame({"a": a, "b": a}) d = pd.Index(b) - check(assert_type(pdutil.hash_pandas_object(b), pd.Series), pd.Series) - check(assert_type(pdutil.hash_pandas_object(c), pd.Series), pd.Series) - check(assert_type(pdutil.hash_pandas_object(d), pd.Series), pd.Series) + check( + assert_type(pdutil.hash_pandas_object(b), "pd.Series[int]"), + pd.Series, + np.uint64, + ) + check( + assert_type(pdutil.hash_pandas_object(c), "pd.Series[int]"), + pd.Series, + np.uint64, + ) + check( + assert_type(pdutil.hash_pandas_object(d), "pd.Series[int]"), + pd.Series, + np.uint64, + ) check( assert_type( pdutil.hash_pandas_object( d, index=True, encoding="latin1", hash_key="apple", categorize=True ), - pd.Series, + "pd.Series[int]", ), pd.Series, + np.uint64, ) @@ -733,7 +706,7 @@ def test_eval() -> None: check( assert_type( pd.eval("double_age = df.age * 2", target=df), - npt.NDArray | Scalar | pd.DataFrame | pd.Series | None, + Scalar | np_ndarray | pd.DataFrame | pd.Series | None, ), pd.DataFrame, ) @@ -751,61 +724,47 @@ def test_to_numeric_scalar() -> None: def test_to_numeric_array_like() -> None: check( - assert_type( - pd.to_numeric([1, 2, 3]), - npt.NDArray, - ), - np.ndarray, + assert_type(pd.to_numeric([1, 2, 3]), np_1darray_anyint), np_1darray, np.integer ) check( - assert_type( - pd.to_numeric([1.0, 2.0, 3.0]), - npt.NDArray, - ), - np.ndarray, + assert_type(pd.to_numeric([1.0, 2.0, 3.0]), np_1darray_float), + np_1darray, + np.floating, ) check( - assert_type( - pd.to_numeric([1.0, 2.0, "3.0"]), - npt.NDArray, - ), - np.ndarray, + assert_type(pd.to_numeric([1.0, 2.0, "3.0"]), np_1darray_float), + np_1darray, + np.floating, ) check( assert_type( pd.to_numeric(np.array([1.0, 2.0, "3.0"], dtype=object)), - npt.NDArray, + np_1darray_float, ), - np.ndarray, + np_1darray, + np.floating, ) check( assert_type( - pd.to_numeric([1.0, 2.0, "blerg"], errors="coerce"), - npt.NDArray, + pd.to_numeric([1.0, 2.0, "blerg"], errors="coerce"), np_1darray_float ), - np.ndarray, + np_1darray, + np.floating, ) check( - assert_type( - pd.to_numeric((1.0, 2.0, 3.0)), - npt.NDArray, - ), - np.ndarray, + assert_type(pd.to_numeric((1.0, 2.0, 3.0)), np_1darray_float), + np_1darray, + np.floating, ) check( - assert_type(pd.to_numeric([1, 2, 3], downcast="unsigned"), npt.NDArray), - np.ndarray, + assert_type(pd.to_numeric([1, 2, 3], downcast="unsigned"), np_1darray_anyint), + np_1darray, + np.integer, ) def test_to_numeric_array_series() -> None: - check( - assert_type( - pd.to_numeric(pd.Series([1, 2, 3])), - pd.Series, - ), - pd.Series, - ) + check(assert_type(pd.to_numeric(pd.Series([1, 2, 3])), pd.Series), pd.Series) check( assert_type( pd.to_numeric(pd.Series([1, 2, "blerg"]), errors="coerce"), @@ -941,41 +900,43 @@ def test_lreshape() -> None: def test_factorize() -> None: codes, uniques = pd.factorize(np.array(["b", "b", "a", "c", "b"])) - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(uniques, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) + check(assert_type(uniques, np_1darray), np_1darray) codes, uniques = pd.factorize(np.recarray((1,), dtype=[("x", int)])) - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(uniques, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) + check(assert_type(uniques, np_1darray), np_1darray) codes, cat_uniques = pd.factorize(pd.Categorical(["b", "b", "a", "c", "b"])) - check(assert_type(codes, np_1darray), np_1darray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) 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_1darray), np_1darray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) 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_1darray), np_1darray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) check(assert_type(idx_uniques, pd.Index), pd.Index) codes, uniques = pd.factorize(np.array(list("bbacb"))) - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(uniques, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) + check(assert_type(uniques, np_1darray), np_1darray) codes, uniques = pd.factorize( np.array(["b", "b", "a", "c", "b"]), use_na_sentinel=True, size_hint=10 ) - check(assert_type(codes, np.ndarray), np.ndarray) - check(assert_type(uniques, np.ndarray), np.ndarray) + check(assert_type(codes, np_1darray[np.int64]), np_1darray[np.int64]) + check(assert_type(uniques, np_1darray), np_1darray) def test_index_unqiue() -> None: ci = pd.CategoricalIndex(["a", "b", "a", "c"]) dti = pd.DatetimeIndex([pd.Timestamp(2000, 1, 1)]) + dti_zoned = pd.Index([dt.tz_localize("Europe/Prague") for dt in dti]) i = pd.Index(["a", "b", "c", "a"]) + ii_pd = pd.Index([1, 2, 3, 3], dtype="Int64") pi = pd.period_range("2000Q1", periods=2, freq="Q") ri = pd.RangeIndex(0, 10) @@ -986,18 +947,35 @@ def test_index_unqiue() -> None: check(assert_type(pd.unique(ci), pd.CategoricalIndex), pd.CategoricalIndex) check( - assert_type(pd.unique(dti), np.ndarray), np.ndarray if PD_LTE_23 else pd.Index + assert_type(pd.unique(dti), np_1darray_dt | pd.DatetimeIndex), + np_1darray if PD_LTE_23 else pd.DatetimeIndex, ) - check(assert_type(pd.unique(i), np.ndarray), np.ndarray if PD_LTE_23 else pd.Index) + check( + assert_type(pd.unique(dti_zoned), np_1darray_dt | pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type(pd.unique(i), np_1darray | pd.Index), + np_1darray if PD_LTE_23 else pd.Index, + ) + check(assert_type(pd.unique(ii_pd), np_1darray | pd.Index), pd.Index) check(assert_type(pd.unique(pi), pd.PeriodIndex), pd.PeriodIndex) - check(assert_type(pd.unique(ri), np.ndarray), np.ndarray if PD_LTE_23 else pd.Index) check( - assert_type(pd.unique(tdi), np.ndarray), np.ndarray if PD_LTE_23 else pd.Index + assert_type(pd.unique(ri), np_1darray_int64), + np_1darray_int64 if PD_LTE_23 else pd.Index, + ) + check( + assert_type(pd.unique(tdi), np_1darray_td), + np_1darray if PD_LTE_23 else pd.Index, + ) + check( + assert_type(pd.unique(mi), np_ndarray), + np_ndarray if PD_LTE_23 else pd.MultiIndex, ) - check(assert_type(pd.unique(mi), np.ndarray), np.ndarray if PD_LTE_23 else pd.Index) check( assert_type(pd.unique(interval_i), "pd.IntervalIndex[pd.Interval[int]]"), pd.IntervalIndex, + pd.Interval, ) @@ -1007,7 +985,10 @@ def test_cut() -> None: b = pd.cut([1, 2, 3, 4, 5, 6, 7, 8], 4, labels=False, duplicates="raise") c = pd.cut([1, 2, 3, 4, 5, 6, 7, 8], 4, labels=["1", "2", "3", "4"]) check(assert_type(a, pd.Categorical), pd.Categorical) - check(assert_type(b, npt.NDArray[np.intp]), np.ndarray) + check( + assert_type(b, np_1darray[np.intp]), + np_1darray[np.intp], + ) check(assert_type(c, pd.Categorical), pd.Categorical) d0, d1 = pd.cut([1, 2, 3, 4, 5, 6, 7, 8], 4, retbins=True) @@ -1016,11 +997,11 @@ def test_cut() -> None: [1, 2, 3, 4, 5, 6, 7, 8], 4, labels=["1", "2", "3", "4"], retbins=True ) check(assert_type(d0, pd.Categorical), pd.Categorical) - check(assert_type(d1, npt.NDArray), np.ndarray) - check(assert_type(e0, npt.NDArray[np.intp]), np.ndarray) - check(assert_type(e1, npt.NDArray), np.ndarray) + check(assert_type(d1, np_1darray_float), np_1darray, np.floating) + check(assert_type(e0, np_1darray[np.intp]), np_1darray[np.intp]) + check(assert_type(e1, np_1darray_float), np_1darray, np.floating) check(assert_type(f0, pd.Categorical), pd.Categorical) - check(assert_type(f1, npt.NDArray), np.ndarray) + check(assert_type(f1, np_1darray_float), np_1darray, np.floating) g = pd.cut(pd.Series([1, 2, 3, 4, 5, 6, 7, 8]), 4, precision=1, duplicates="drop") h = pd.cut(pd.Series([1, 2, 3, 4, 5, 6, 7, 8]), 4, labels=False, duplicates="raise") @@ -1055,11 +1036,11 @@ def test_cut() -> None: retbins=True, ) check(assert_type(j0, pd.Series), pd.Series) - check(assert_type(j1, npt.NDArray), np.ndarray) + check(assert_type(j1, np_1darray_float), np_1darray, np.floating) check(assert_type(k0, pd.Series), pd.Series) - check(assert_type(k1, npt.NDArray), np.ndarray) + check(assert_type(k1, np_1darray_float), np_1darray, np.floating) check(assert_type(l0, pd.Series), pd.Series) - check(assert_type(l1, npt.NDArray), np.ndarray) + check(assert_type(l1, np_1darray_float), np_1darray, np.floating) check(assert_type(m0, pd.Series), pd.Series) check(assert_type(m1, pd.IntervalIndex), pd.IntervalIndex) @@ -1132,10 +1113,10 @@ def test_qcut() -> None: check(assert_type(c0, pd.Categorical), pd.Categorical) check(assert_type(d0, pd.Series), pd.Series) - check(assert_type(a1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(b1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(c1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(d1, npt.NDArray[np.double]), np.ndarray) + check(assert_type(a1, np_1darray_float), np_1darray, np.floating) + check(assert_type(b1, np_1darray_float), np_1darray, np.floating) + check(assert_type(c1, np_1darray_float), np_1darray, np.floating) + check(assert_type(d1, np_1darray_float), np_1darray, np.floating) e0, e1 = pd.qcut(val_list, [0.25, 0.5, 0.75], retbins=True) f0, f1 = pd.qcut(val_arr, np.array([0.25, 0.5, 0.75]), retbins=True) @@ -1146,17 +1127,18 @@ def test_qcut() -> None: check(assert_type(e0, pd.Categorical), pd.Categorical) check(assert_type(f0, pd.Categorical), pd.Categorical) - check(assert_type(g0, npt.NDArray[np.intp]), np.ndarray) + check(assert_type(g0, np_1darray_intp | np_1darray_float), np_1darray_intp) check(assert_type(h0, pd.Series), pd.Series) - check(assert_type(i0, npt.NDArray[np.intp]), np.ndarray) - check(assert_type(j0, npt.NDArray[np.intp]), np.ndarray) - - check(assert_type(e1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(f1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(g1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(h1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(i1, npt.NDArray[np.double]), np.ndarray) - check(assert_type(j1, npt.NDArray[np.double]), np.ndarray) + # because of nans + check(assert_type(i0, np_1darray_intp | np_1darray_float), np_1darray, np.floating) + check(assert_type(j0, np_1darray_intp | np_1darray_float), np_1darray, np.floating) + + check(assert_type(e1, np_1darray_float), np_1darray, np.floating) + check(assert_type(f1, np_1darray_float), np_1darray, np.floating) + check(assert_type(g1, np_1darray_float), np_1darray, np.floating) + check(assert_type(h1, np_1darray_float), np_1darray, np.floating) + check(assert_type(i1, np_1darray_float), np_1darray, np.floating) + check(assert_type(j1, np_1darray_float), np_1darray, np.floating) @pytest.mark.xfail( diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 8565f4753..696e57f8b 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -380,7 +380,7 @@ def test_scatter_matrix(close_figures: None) -> None: pd.plotting.scatter_matrix(df, alpha=0.2), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -399,14 +399,14 @@ def test_plot_line() -> None: IRIS_DF.plot.line(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="line", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -418,14 +418,14 @@ def test_plot_area(close_figures: None) -> None: IRIS_DF.plot.area(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="area", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -437,14 +437,14 @@ def test_plot_bar(close_figures: None) -> None: IRIS_DF.plot.bar(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="bar", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -456,14 +456,14 @@ def test_plot_barh(close_figures: None) -> None: IRIS_DF.plot.barh(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="barh", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -494,14 +494,14 @@ def test_plot_density(close_figures: None) -> None: IRIS_DF.plot.density(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="density", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -519,14 +519,14 @@ def test_plot_hexbin(close_figures: None) -> None: IRIS_DF.plot.hexbin(x="SepalLength", y="SepalWidth", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(kind="hexbin", x="SepalLength", y="SepalWidth", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -538,14 +538,14 @@ def test_plot_hist(close_figures: None) -> None: IRIS_DF.plot.hist(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(subplots=True, kind="hist"), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -557,14 +557,14 @@ def test_plot_kde(close_figures: None) -> None: IRIS_DF.plot.kde(subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( IRIS_DF.plot(subplots=True, kind="kde"), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -576,7 +576,7 @@ def test_plot_pie(close_figures: None) -> None: IRIS_DF.plot.pie(y="SepalLength", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( @@ -584,7 +584,7 @@ def test_plot_pie(close_figures: None) -> None: IRIS_DF.plot(kind="pie", y="SepalLength", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -604,7 +604,7 @@ def test_plot_scatter(close_figures: None) -> None: IRIS_DF.plot.scatter(x="SepalLength", y="SepalWidth", subplots=True), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) check( assert_type( @@ -613,7 +613,7 @@ def test_plot_scatter(close_figures: None) -> None: ), npt.NDArray[np.object_], ), - np.ndarray, + npt.NDArray[np.object_], ) @@ -697,7 +697,7 @@ def test_plot_subplot_changes_150() -> None: assert_type( df.plot(subplots=[("a", "b"), ("c", "d")]), npt.NDArray[np.object_] ), - np.ndarray, + npt.NDArray[np.object_], ) diff --git a/tests/test_styler.py b/tests/test_styler.py index 16f839020..bbaf44bf7 100644 --- a/tests/test_styler.py +++ b/tests/test_styler.py @@ -9,7 +9,6 @@ ) from jinja2.loaders import PackageLoader import numpy as np -import numpy.typing as npt from pandas import ( DataFrame, Index, @@ -24,6 +23,7 @@ PD_LTE_23, check, ensure_clean, + np_ndarray_str, ) from pandas.io.formats.style import Styler @@ -64,7 +64,7 @@ def highlight_max(x: Series[int], /, color: str) -> list[str]: def test_apply_index() -> None: - def f(s: Series) -> npt.NDArray[np.str_]: + def f(s: Series) -> np_ndarray_str: return np.asarray(s, dtype=np.str_) check(assert_type(DF.style.apply_index(f), Styler), Styler) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 7753f9850..e495209f7 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -29,6 +29,9 @@ TYPE_CHECKING_INVALID_USAGE, check, np_1darray, + np_1darray_dt, + np_1darray_int64, + np_1darray_td, pytest_warns_bounded, ) @@ -892,27 +895,17 @@ def test_series_types_to_numpy() -> None: i_s = pd.interval_range(1, 2).to_series() # default dtype + check(assert_type(td_s.to_numpy(), np_1darray_td), np_1darray, np.timedelta64) check( - assert_type(td_s.to_numpy(), np_1darray[np.timedelta64]), + assert_type(td_s.to_numpy(na_value=pd.Timedelta(0)), np_1darray_td), np_1darray, - dtype=np.timedelta64, + np.timedelta64, ) + check(assert_type(ts_s.to_numpy(), np_1darray_dt), np_1darray, np.datetime64) 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]), + assert_type(ts_s.to_numpy(na_value=pd.Timestamp(1)), np_1darray_dt), np_1darray, - dtype=np.datetime64, + np.datetime64, ) check( assert_type(p_s.to_numpy(), np_1darray[np.object_]), @@ -978,29 +971,29 @@ def test_series_types_to_numpy() -> None: # passed dtype-like with statically known generic check( - assert_type(td_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(td_s.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( - assert_type(td_s.to_numpy(dtype=np.timedelta64), np_1darray[np.timedelta64]), + assert_type(td_s.to_numpy(dtype=np.timedelta64), np_1darray_td), np_1darray, - dtype=np.timedelta64, + np.timedelta64, ) check( - assert_type(ts_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(ts_s.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( - assert_type(ts_s.to_numpy(dtype=np.datetime64), np_1darray[np.datetime64]), + assert_type(ts_s.to_numpy(dtype=np.datetime64), np_1darray_dt), np_1darray, - dtype=np.datetime64, + np.datetime64, ) check( - assert_type(p_s.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(p_s.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( assert_type(o_s.to_numpy(dtype=np.bytes_), np_1darray[np.bytes_]), @@ -1026,27 +1019,17 @@ def test_index_types_to_numpy() -> None: i_i = pd.interval_range(1, 2) # default dtype + check(assert_type(td_i.to_numpy(), np_1darray_td), np_1darray, np.timedelta64) 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]), + assert_type(td_i.to_numpy(na_value=pd.Timedelta(0)), np_1darray_td), np_1darray, - dtype=np.datetime64, + np.timedelta64, ) + check(assert_type(ts_i.to_numpy(), np_1darray_dt), np_1darray, np.datetime64) check( - assert_type(ts_i.to_numpy(na_value=pd.Timestamp(1)), np_1darray[np.datetime64]), + assert_type(ts_i.to_numpy(na_value=pd.Timestamp(1)), np_1darray_dt), np_1darray, - dtype=np.datetime64, + np.datetime64, ) check( assert_type(p_i.to_numpy(), np_1darray[np.object_]), @@ -1095,19 +1078,19 @@ def test_index_types_to_numpy() -> None: # passed dtype-like with statically known generic check( - assert_type(td_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(td_i.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( - assert_type(ts_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(ts_i.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( - assert_type(p_i.to_numpy(dtype=np.int64), np_1darray[np.int64]), + assert_type(p_i.to_numpy(dtype=np.int64), np_1darray_int64), np_1darray, - dtype=np.int64, + np.int64, ) check( assert_type(i_i.to_numpy(dtype=np.bytes_), np_1darray[np.bytes_]), diff --git a/tests/test_windowing.py b/tests/test_windowing.py index 8abd91f9d..7e302b449 100644 --- a/tests/test_windowing.py +++ b/tests/test_windowing.py @@ -8,6 +8,10 @@ Timedelta, date_range, ) +from pandas.core.indexers.objects import ( + FixedForwardWindowIndexer, + VariableOffsetWindowIndexer, +) from pandas.core.window import ( Rolling, Window, @@ -19,6 +23,8 @@ from tests import ( PD_LTE_23, check, + np_1darray_intp, + np_ndarray, pytest_warns_bounded, ) @@ -88,7 +94,7 @@ def _mean(df: DataFrame) -> Series: check(assert_type(DF.rolling(10).apply(_mean), DataFrame), DataFrame) - def _mean2(df: DataFrame) -> np.ndarray: + def _mean2(df: DataFrame) -> np_ndarray: return np.mean(df, axis=0) check(assert_type(DF.rolling(10).apply(_mean2, raw=True), DataFrame), DataFrame) @@ -181,7 +187,7 @@ def _mean(df: Series) -> float: check(assert_type(S.rolling(10).apply(_mean), Series), Series) - def _mean2(df: Series) -> np.ndarray: + def _mean2(df: Series) -> np_ndarray: return np.mean(df, axis=0) check(assert_type(S.rolling(10).apply(_mean2, raw=True), Series), Series) @@ -247,7 +253,7 @@ def _mean(df: DataFrame) -> Series: check(assert_type(DF.expanding(10).apply(_mean), DataFrame), DataFrame) - def _mean2(df: DataFrame) -> np.ndarray: + def _mean2(df: DataFrame) -> np_ndarray: return np.mean(df, axis=0) check(assert_type(DF.expanding(10).apply(_mean2, raw=True), DataFrame), DataFrame) @@ -306,7 +312,7 @@ def _mean(df: Series) -> float: check(assert_type(S.expanding(10).apply(_mean), Series), Series) - def _mean2(df: Series) -> np.ndarray: + def _mean2(df: Series) -> np_ndarray: return np.mean(df, axis=0) check(assert_type(S.expanding(10).apply(_mean2, raw=True), Series), Series) @@ -414,7 +420,7 @@ def test_rolling_window() -> None: ], ) - indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=2) + indexer = FixedForwardWindowIndexer(window_size=2) check( assert_type(df_time.rolling(window=indexer, min_periods=1).sum(), DataFrame), DataFrame, @@ -424,3 +430,27 @@ def test_rolling_window() -> None: assert_type(s.rolling(window=indexer, min_periods=1).sum(), Series), Series, ) + + +def test_indexer_variable_offset() -> None: + indexer = VariableOffsetWindowIndexer(index=IDX, offset=pd.offsets.BDay(1)) + check( + assert_type( + indexer.get_window_bounds(2, None, None), + tuple[np_1darray_intp, np_1darray_intp], + ), + tuple, + np.ndarray, + ) + + +def test_indexer_fixed_forward() -> None: + indexer = FixedForwardWindowIndexer([1, float("nan")], 2) + check( + assert_type( + indexer.get_window_bounds(2, None, None), + tuple[np_1darray_intp, np_1darray_intp], + ), + tuple, + np.ndarray, + )