diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index 5152f32b9..9aaede454 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -4,6 +4,7 @@ from typing import ( Literal, TypeVar, overload, + type_check_only, ) import numpy as np @@ -18,7 +19,6 @@ from pandas._typing import ( IntervalClosedType, IntervalT, np_1darray, - npt, ) 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__( @@ -36,9 +37,8 @@ class _LengthDescriptor: def __get__( self, instance: Interval[_OrderableTimesT], owner: Any ) -> Timedelta: ... - @overload - def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... +@type_check_only class _MidDescriptor: @overload def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ... @@ -46,8 +46,6 @@ class _MidDescriptor: def __get__( self, instance: Interval[_OrderableTimesT], owner: Any ) -> _OrderableTimesT: ... - @overload - def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... class IntervalMixin: @property @@ -68,8 +66,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, @@ -223,21 +221,4 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @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: ... +class IntervalTree(IntervalMixin): ... 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/_testing/__init__.pyi b/pandas-stubs/_testing/__init__.pyi index 3db4eccbb..06d605cdf 100644 --- a/pandas-stubs/_testing/__init__.pyi +++ b/pandas-stubs/_testing/__init__.pyi @@ -12,7 +12,7 @@ from typing import ( import warnings from matplotlib.artist import Artist -import numpy as np +from numpy import typing as npt from pandas import ( Categorical, DataFrame, @@ -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 | npt.NDArray[Any] | 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 | npt.NDArray[Any] | 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 f1e2bf606..03fa4b426 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"] @@ -832,18 +833,6 @@ 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) @@ -871,6 +860,17 @@ np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]] np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]] 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[Any] | tuple[Any, ...] | Index +) +ListLikeU: TypeAlias = Sequence | np_1darray[Any] | Series | Index +ListLikeHashable: TypeAlias = ( + MutableSequence[HashableT] | np_1darray[Any] | 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..2b3b5c867 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -5,65 +5,75 @@ from typing import ( ) 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.period import PeriodIndex +from pandas.core.series import Series from pandas._typing import ( + T_EXTENSION_ARRAY, AnyArrayLike, + GenericT, IntervalT, TakeIndexer, np_1darray, + 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 +def unique(values: DatetimeIndex) -> np_1darray[np.datetime64] | DatetimeIndex: ... +@overload +def unique(values: Index) -> np_1darray[Any] | Index: ... @overload def unique(values: Categorical) -> Categorical: ... + +# @overload +# def unique(values: Series[Never]) -> np_1darray[Any] | ExtensionArray: ... +# TODO: DatetimeArray python/mypy#19952 +# @overload +# def unique(values: Series[Timestamp]) -> np_1darray[np.datetime64] | ExtensionArray: ... +# @overload +# def unique(values: Series[int]) -> np_1darray[np.integer] | ExtensionArray: ... @overload -def unique(values: Series) -> np.ndarray | ExtensionArray: ... +def unique(values: Series) -> np_1darray[Any] | 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[np.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[np.int64], Index]: ... @overload def factorize( values: Categorical, sort: bool = ..., use_na_sentinel: bool = ..., size_hint: int | None = ..., -) -> tuple[np_1darray, Categorical]: ... +) -> tuple[np_1darray[np.int64], Categorical]: ... def value_counts( values: AnyArrayLike | list | tuple, sort: bool = True, @@ -73,9 +83,9 @@ 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, fill_value: Any = None, -) -> np_1darray | ExtensionArray: ... +) -> np_1darray[Any] | ExtensionArray: ... diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index 09bf2b0cb..e29153bed 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -35,7 +35,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 | npt.NDArray[Any], value: Any) -> None: ... def __len__(self) -> int: ... def __iter__(self) -> Iterator[Any]: ... def __contains__(self, item: object) -> bool | np.bool_: ... @@ -83,7 +83,9 @@ class ExtensionArray: side: Literal["left", "right"] = ..., sorter: ListLike | None = ..., ) -> np.intp: ... - def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ... + def factorize( + self, use_na_sentinel: bool = True + ) -> tuple[np_1darray[Any], Self]: ... def repeat( self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None ) -> 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/datetimelike.pyi b/pandas-stubs/core/arrays/datetimelike.pyi index 1f4870898..d764e1ea4 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,7 +64,7 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, ExtensionArray): def ravel(self, *args: Any, **kwargs: Any): ... def __iter__(self): ... @property - def asi8(self) -> np.ndarray: ... + def asi8(self) -> np_1darray[Any]: ... @property def nbytes(self): ... def __array__( diff --git a/pandas-stubs/core/arrays/interval.pyi b/pandas-stubs/core/arrays/interval.pyi index 52d23ac51..cfa2a570c 100644 --- a/pandas-stubs/core/arrays/interval.pyi +++ b/pandas-stubs/core/arrays/interval.pyi @@ -5,6 +5,7 @@ from typing import ( ) import numpy as np +from numpy import typing as npt from pandas.core.arrays.base import ExtensionArray as ExtensionArray from pandas.core.indexes.base import Index from pandas.core.series import Series @@ -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 + self, other: Scalar | ExtensionArray | Index | npt.NDArray[Any] ) -> np_1darray[np.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..b9a2d5139 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[Any]: ... __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 ea3865add..5a3f4d709 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -151,7 +151,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]): def is_monotonic_increasing(self) -> bool: ... def factorize( self, sort: bool = False, use_na_sentinel: bool = True - ) -> tuple[np_1darray, np_1darray | Index | Categorical]: ... + ) -> tuple[np_1darray[Any], np_1darray | Index | Categorical]: ... @overload def searchsorted( self, @@ -174,7 +174,7 @@ NumListLike: TypeAlias = ( | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex - | dict[str, np.ndarray] + | dict[str, np_1darray[Any]] | Sequence[complex] ) diff --git a/pandas-stubs/core/computation/eval.pyi b/pandas-stubs/core/computation/eval.pyi index 19df29909..a6927e1ef 100644 --- a/pandas-stubs/core/computation/eval.pyi +++ b/pandas-stubs/core/computation/eval.pyi @@ -25,4 +25,4 @@ def eval( level: int = 0, target: object | None = None, inplace: bool = False, -) -> npt.NDArray | Scalar | DataFrame | Series | None: ... +) -> npt.NDArray[Any] | Scalar | DataFrame | Series | None: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 93af6f79d..ec777e851 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, @@ -162,8 +163,7 @@ from pandas._typing import ( WriteBuffer, XMLParsers, np_2darray, - npt, - num, + np_ndarray_float, ) from pandas.io.formats.style import Styler @@ -206,7 +206,7 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]): Scalar | Series | DataFrame - | np.ndarray + | npt.NDArray[Any] | NAType | NaTType | Mapping[Hashable, Scalar | NAType | NaTType] @@ -427,7 +427,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def __matmul__(self, other: Series) -> Series: ... @overload - def __matmul__(self, other: np.ndarray) -> Self: ... + def __matmul__(self, other: npt.NDArray[Any]) -> Self: ... def __rmatmul__(self, other): ... @overload @classmethod @@ -832,7 +832,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def eval( self, expr: _str, *, inplace: Literal[False] = False, **kwargs: Any - ) -> Scalar | np.ndarray | Self | Series: ... + ) -> Scalar | npt.NDArray[Any] | Self | Series: ... @overload def select_dtypes( self, include: StrDtypeArg, exclude: _AstypeArgExtList | None = ... @@ -970,7 +970,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): Label | Series | Index - | np.ndarray + | npt.NDArray[Any] | Iterator[Hashable] | Sequence[Hashable] ), @@ -987,7 +987,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): Label | Series | Index - | np.ndarray + | npt.NDArray[Any] | Iterator[Hashable] | Sequence[Hashable] ), @@ -1402,8 +1402,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 | npt.NDArray[Any] | None = ..., + value_vars: tuple | Sequence | npt.NDArray[Any] | None = ..., var_name: Scalar | None = None, value_name: Scalar = "value", col_level: int | _str | None = ..., @@ -1676,7 +1676,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 = ..., @@ -1778,7 +1778,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, @@ -1786,7 +1786,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, @@ -1794,7 +1794,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, @@ -1802,7 +1802,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, @@ -2005,20 +2005,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: ... @@ -2074,7 +2068,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, @@ -2144,7 +2138,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | npt.NDArray[Any] | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2160,7 +2154,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | npt.NDArray[Any] | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2200,21 +2194,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, @@ -2232,7 +2219,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, @@ -2554,13 +2541,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, @@ -2601,7 +2581,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | npt.NDArray[Any] | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), @@ -2617,7 +2597,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): cond: ( Series | DataFrame - | np.ndarray + | npt.NDArray[Any] | Callable[[DataFrame], DataFrame] | Callable[[Any], _bool] ), diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 1b292b843..48402be7c 100644 --- a/pandas-stubs/core/generic.pyi +++ b/pandas-stubs/core/generic.pyi @@ -21,7 +21,7 @@ from typing import ( overload, ) -import numpy as np +from numpy import typing as npt from pandas import Index import pandas.core.indexing as indexing from pandas.core.resample import DatetimeIndexResampler @@ -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=...) -> npt.NDArray[Any]: ... @final def to_excel( self, diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 7c058cd79..3279e31b0 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -19,7 +19,7 @@ from typing import ( ) from matplotlib.axes import Axes as PlotAxes -import numpy as np +from numpy import typing as npt from pandas.core.frame import DataFrame from pandas.core.groupby.base import TransformReductionListType from pandas.core.groupby.groupby import ( @@ -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[[npt.NDArray[Any], npt.NDArray[Any]], 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..abb89e22b 100644 --- a/pandas-stubs/core/groupby/groupby.pyi +++ b/pandas-stubs/core/groupby/groupby.pyi @@ -240,7 +240,7 @@ class GroupBy(BaseGroupBy[NDFrameT]): adjust: bool = ..., ignore_na: bool = ..., axis: Axis = ..., - times: str | np.ndarray | Series | np.timedelta64 | None = ..., + times: str | npt.NDArray[Any] | Series | np.timedelta64 | None = ..., method: CalculationMethod = ..., *, selection: IndexLabel | None = ..., diff --git a/pandas-stubs/core/groupby/grouper.pyi b/pandas-stubs/core/groupby/grouper.pyi index 8ac7f901a..30900f48a 100644 --- a/pandas-stubs/core/groupby/grouper.pyi +++ b/pandas-stubs/core/groupby/grouper.pyi @@ -30,6 +30,8 @@ from pandas._typing import ( ) from pandas.util._decorators import cache_readonly +from pandas.io.formats.printing import PrettyDict + class Grouper: key: KeysArgType | None level: Level | ListLikeHashable[Level] | None @@ -73,3 +75,5 @@ class Grouping: def result_index(self) -> Index: ... @cache_readonly def group_index(self) -> Index: ... + @cache_readonly + def groups(self) -> PrettyDict[Hashable, Index]: ... diff --git a/pandas-stubs/core/groupby/ops.pyi b/pandas-stubs/core/groupby/ops.pyi index 740eeb760..9a77edc6d 100644 --- a/pandas-stubs/core/groupby/ops.pyi +++ b/pandas-stubs/core/groupby/ops.pyi @@ -53,7 +53,7 @@ class BaseGrouper: @final def size(self) -> Series: ... @cache_readonly - def groups(self) -> dict[Hashable, np.ndarray]: ... + def groups(self) -> dict[Hashable, Index]: ... @final @cache_readonly def is_monotonic(self) -> bool: ... diff --git a/pandas-stubs/core/indexers.pyi b/pandas-stubs/core/indexers.pyi index c6bcc6d84..6d5db972d 100644 --- a/pandas-stubs/core/indexers.pyi +++ b/pandas-stubs/core/indexers.pyi @@ -1,13 +1,21 @@ from typing import Any import numpy as np +from numpy import typing as npt -def check_array_indexer(arrayArrayLike, indexer): ... +from pandas._typing import ( + AnyArrayLike, + np_1darray, +) + +def check_array_indexer( + arrayArrayLike: AnyArrayLike, indexer: AnyArrayLike +) -> np_1darray[np.bool_]: ... class BaseIndexer: def __init__( self, - index_array: np.ndarray | None = ..., + index_array: npt.NDArray[Any] | None = ..., window_size: int = ..., **kwargs: Any, ) -> None: ... @@ -17,12 +25,12 @@ class BaseIndexer: min_periods: int | None = ..., center: bool | None = ..., closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... + ) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: ... class VariableOffsetWindowIndexer(BaseIndexer): def __init__( self, - index_array: np.ndarray | None = ..., + index_array: npt.NDArray[Any] | None = ..., window_size: int = ..., index=..., offset=..., @@ -34,7 +42,7 @@ class VariableOffsetWindowIndexer(BaseIndexer): min_periods: int | None = ..., center: bool | None = ..., closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... + ) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: ... class FixedForwardWindowIndexer(BaseIndexer): def get_window_bounds( @@ -43,4 +51,4 @@ class FixedForwardWindowIndexer(BaseIndexer): min_periods: int | None = ..., center: bool | None = ..., closed: str | None = ..., - ) -> tuple[np.ndarray, np.ndarray]: ... + ) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: ... diff --git a/pandas-stubs/core/indexers/objects.pyi b/pandas-stubs/core/indexers/objects.pyi index 2187cf885..74541b805 100644 --- a/pandas-stubs/core/indexers/objects.pyi +++ b/pandas-stubs/core/indexers/objects.pyi @@ -1,13 +1,16 @@ from typing import Any -import numpy as np -from pandas import DatetimeIndex +from numpy import typing as npt +from pandas.core.indexes.datetimes import DatetimeIndex from pandas._libs.tslibs import BaseOffset class BaseIndexer: def __init__( - self, index_array: np.ndarray | None = None, window_size: int = 0, **kwargs: Any + self, + index_array: npt.NDArray[Any] | None = None, + window_size: int = 0, + **kwargs: Any, ) -> None: ... def get_window_bounds( self, @@ -16,14 +19,14 @@ class BaseIndexer: center: bool | None, closed: str | None = None, step: int | None = None, - ) -> tuple[np.ndarray, np.ndarray]: ... + ) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: ... class FixedForwardWindowIndexer(BaseIndexer): ... class VariableOffsetWindowIndexer(BaseIndexer): def __init__( self, - index_array: np.ndarray | None = None, + index_array: npt.NDArray[Any] | None = None, window_size: int = 0, index: DatetimeIndex | None = None, offset: BaseOffset | None = None, diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 8252a0228..318c4689c 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -106,6 +106,8 @@ from pandas._typing import ( type_t, ) +from pandas.io.formats.printing import PrettyDict + class InvalidIndexError(Exception): ... class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @@ -469,7 +471,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): sort: bool = False, ) -> Index: ... @property - def values(self) -> np_1darray: ... + def values(self) -> np_1darray[Any]: ... def memory_usage(self, deep: bool = False): ... @overload def where( @@ -521,7 +523,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @final def get_indexer_for(self, target, **kwargs: Any): ... @final - def groupby(self, values) -> dict[Hashable, np.ndarray]: ... + def groupby(self, values) -> PrettyDict[Hashable, Self]: ... def map(self, mapper, na_action=...) -> Index: ... def isin(self, values, level=...) -> np_1darray[np.bool]: ... def slice_indexer( diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index dd9f12471..d2112cb3e 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/pandas-stubs/core/indexes/category.pyi @@ -2,9 +2,12 @@ from collections.abc import ( Hashable, Iterable, ) -from typing import final +from typing import ( + Any, + final, +) -import numpy as np +from numpy import typing as npt from pandas.core import accessor from pandas.core.arrays.categorical import Categorical from pandas.core.indexes.base import Index @@ -14,7 +17,7 @@ from typing_extensions import Self from pandas._typing import S1 class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): - codes: np.ndarray = ... + codes: npt.NDArray[Any] = ... categories: Index = ... @property def array(self) -> Categorical: ... # type: ignore[override] # pyrefly: ignore[bad-override] diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 84abc62ea..4ae2a9457 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/pandas-stubs/core/indexes/period.pyi @@ -6,6 +6,7 @@ from typing import ( ) import numpy as np +from numpy import typing as npt import pandas as pd from pandas import Index from pandas.core.indexes.accessors import PeriodIndexFieldOps @@ -67,7 +68,7 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexField def asof_locs( self, where: pd.DatetimeIndex | PeriodIndex, - mask: np_1darray[np.bool_], + mask: npt.NDArray[np.bool_], ) -> np_1darray[np.intp]: ... @property def is_full(self) -> bool: ... diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index dd1a8b189..30df91413 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -36,7 +36,6 @@ from pandas._typing import ( np_ndarray_dt, np_ndarray_float, np_ndarray_td, - num, ) _NUM_FACTOR: TypeAlias = Just[int] | Just[float] | np.integer | np.floating @@ -125,7 +124,7 @@ class TimedeltaIndex( ) -> Index[float]: ... @overload # type: ignore[override] # pyrefly: ignore # bad-override - def __floordiv__(self, other: num | Sequence[float]) -> Self: ... + def __floordiv__(self, other: float | Sequence[float]) -> Self: ... @overload def __floordiv__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: dt.timedelta | Sequence[dt.timedelta] diff --git a/pandas-stubs/core/resample.pyi b/pandas-stubs/core/resample.pyi index 91299803e..6efd26b8a 100644 --- a/pandas-stubs/core/resample.pyi +++ b/pandas-stubs/core/resample.pyi @@ -12,16 +12,15 @@ from typing import ( ) import numpy as np -from pandas import ( - DataFrame, - Series, - Timedelta, -) +from numpy import typing as npt +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 +29,6 @@ from pandas._typing import ( Scalar, TimeGrouperOrigin, TimestampConvention, - npt, ) _FrameGroupByFunc: TypeAlias = ( diff --git a/pandas-stubs/core/reshape/melt.pyi b/pandas-stubs/core/reshape/melt.pyi index 68a9ee57d..2c9429edd 100644 --- a/pandas-stubs/core/reshape/melt.pyi +++ b/pandas-stubs/core/reshape/melt.pyi @@ -1,14 +1,18 @@ -from collections.abc import Hashable +from collections.abc import ( + Hashable, + Sequence, +) +from typing import Any -import numpy as np +from numpy import typing as npt from pandas.core.frame import DataFrame from pandas._typing import HashableT def melt( frame: DataFrame, - id_vars: tuple | list | np.ndarray | None = None, - value_vars: tuple | list | np.ndarray | None = None, + id_vars: Sequence[Hashable] | npt.NDArray[Any] | None = None, + value_vars: Sequence[Hashable] | npt.NDArray[Any] | 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 8e8382798..9f76c0955 100644 --- a/pandas-stubs/core/reshape/pivot.pyi +++ b/pandas-stubs/core/reshape/pivot.pyi @@ -6,6 +6,7 @@ from collections.abc import ( ) import datetime from typing import ( + Any, Literal, TypeAlias, overload, @@ -85,7 +86,7 @@ def pivot_table( values: _PivotTableValuesTypes = ..., *, index: Grouper, - columns: _PivotTableColumnsTypes | Index | npt.NDArray = ..., + columns: _PivotTableColumnsTypes | Index | npt.NDArray[Any] = ..., aggfunc: ( _PivotAggFunc | Sequence[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc] ) = ..., @@ -100,7 +101,7 @@ def pivot_table( def pivot_table( data: DataFrame, values: _PivotTableValuesTypes = ..., - index: _PivotTableIndexTypes | Index | npt.NDArray = ..., + index: _PivotTableIndexTypes | Index | npt.NDArray[Any] = ..., *, columns: Grouper, aggfunc: ( diff --git a/pandas-stubs/core/reshape/tile.pyi b/pandas-stubs/core/reshape/tile.pyi index 0eed14683..c9a4e3efd 100644 --- a/pandas-stubs/core/reshape/tile.pyi +++ b/pandas-stubs/core/reshape/tile.pyi @@ -1,5 +1,6 @@ from collections.abc import Sequence from typing import ( + Any, Literal, overload, ) @@ -19,12 +20,15 @@ from pandas.core.series import Series from pandas._typing import ( IntervalT, Label, + np_1darray, + np_ndarray_anyint, + np_ndarray_float, npt, ) @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], right: bool = ..., *, @@ -34,10 +38,10 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[npt.NDArray[np.intp], npt.NDArray]: ... +) -> tuple[np_1darray[np.intp], np_1darray[np.double]]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: IntervalIndex[IntervalT], right: bool = ..., *, @@ -47,7 +51,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[npt.NDArray[np.intp], IntervalIndex[IntervalT]]: ... +) -> tuple[np_1darray[np.intp], IntervalIndex[IntervalT]]: ... @overload def cut( # pyright: ignore[reportOverlappingOverload] x: Series[Timestamp], @@ -92,7 +96,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[Series, npt.NDArray]: ... +) -> tuple[Series, np_1darray[np.double]]: ... @overload def cut( x: Series, @@ -108,7 +112,7 @@ def cut( ) -> tuple[Series, IntervalIndex]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], right: bool = ..., labels: Sequence[Label] | None = ..., @@ -118,10 +122,10 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[Categorical, npt.NDArray]: ... +) -> tuple[Categorical, np_1darray[np.double]]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: IntervalIndex[IntervalT], right: bool = ..., labels: Sequence[Label] | None = ..., @@ -134,7 +138,7 @@ def cut( ) -> tuple[Categorical, IntervalIndex[IntervalT]]: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: ( int | Series @@ -152,7 +156,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> npt.NDArray[np.intp]: ... +) -> np_1darray[np.intp]: ... @overload def cut( x: Series[Timestamp], @@ -194,7 +198,7 @@ def cut( ) -> Series: ... @overload def cut( - x: Index | npt.NDArray | Sequence[int] | Sequence[float], + x: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], bins: ( int | Series @@ -214,59 +218,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: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], + 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[np.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: Index | npt.NDArray[Any] | Sequence[int] | Sequence[float], + 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[np.intp | np.double], np_1darray[np.double]]: ... @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[np.double]]: ... @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[np.double]]: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c05af4813..e2a779b2a 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -47,6 +47,7 @@ from matplotlib.axes import ( SubplotBase, ) import numpy as np +from numpy import typing as npt from pandas import ( Index, Period, @@ -60,8 +61,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 ( @@ -210,8 +215,6 @@ from pandas._typing import ( np_ndarray_float, np_ndarray_str, np_ndarray_td, - npt, - num, ) from pandas.core.dtypes.base import ExtensionDtype @@ -296,9 +299,9 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): value: S1 | ArrayLike | Series[S1] | None, ) -> None: ... -_DataLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1] +_DataLike: TypeAlias = ArrayLike | dict[str, npt.NDArray[Any]] | SequenceNotStr[S1] _DataLikeS1: TypeAlias = ( - ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | IndexOpsMixin[S1] + ArrayLike | dict[_str, npt.NDArray[Any]] | Sequence[S1] | IndexOpsMixin[S1] ) class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @@ -510,8 +513,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[Any] | ExtensionArray | Categorical: ... + def ravel(self, order: _str = ...) -> np_1darray[Any]: ... def __len__(self) -> int: ... def view(self, dtype=...) -> Series[S1]: ... @final @@ -806,13 +809,13 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def count(self) -> int: ... def mode(self, dropna=True) -> Series[S1]: ... @overload - def unique(self: Series[Never]) -> np.ndarray: ... # type: ignore[overload-overlap] + def unique(self: Series[Never]) -> np_1darray[Any]: ... # 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[Any]: ... @overload def drop_duplicates( self, @@ -889,20 +892,21 @@ 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, npt.NDArray[Any]] | Sequence[S1] | Index[S1], + ) -> npt.NDArray[Any]: ... @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: npt.NDArray[Any]) -> npt.NDArray[Any]: ... @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: npt.NDArray[Any]) -> npt.NDArray[Any]: ... @overload def searchsorted( self, @@ -1327,7 +1331,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 | npt.NDArray[Any] | None = None, random_state: RandomState | None = None, axis: AxisIndex | None = None, ignore_index: _bool = False, @@ -1546,7 +1550,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | npt.NDArray[Any] | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1562,7 +1566,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | npt.NDArray[Any] | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1578,7 +1582,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | npt.NDArray[Any] | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1594,7 +1598,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): cond: ( Series[S1] | Series[_bool] - | np.ndarray + | npt.NDArray[Any] | Callable[[Series[S1]], Series[bool]] | Callable[[S1], bool] ), @@ -1610,8 +1614,8 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): tuple[ Sequence[bool] | Series[bool] - | Callable[[Series], Series | np.ndarray | Sequence[bool]], - ListLikeU | Scalar | Callable[[Series], Series | np.ndarray], + | Callable[[Series], Series | npt.NDArray[Any] | Sequence[bool]], + ListLikeU | Scalar | Callable[[Series], Series | npt.NDArray[Any]], ], ], ) -> Series: ... @@ -2903,9 +2907,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] # pyrefly: ignore # bad-override @@ -2922,9 +2926,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] # pyrefly: ignore # bad-override @@ -4156,7 +4160,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 = ..., @@ -4179,7 +4183,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: npt.NDArray[Any] | Series | None = None, method: CalculationMethod = "single", ) -> ExponentialMovingWindow[Series]: ... @final @@ -4314,7 +4318,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, @@ -4330,7 +4334,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, diff --git a/pandas-stubs/core/tools/numeric.pyi b/pandas-stubs/core/tools/numeric.pyi index 3e3d100ed..e130f880b 100644 --- a/pandas-stubs/core/tools/numeric.pyi +++ b/pandas-stubs/core/tools/numeric.pyi @@ -1,10 +1,13 @@ +from collections.abc import Sequence from typing import ( + Any, Literal, TypeAlias, overload, ) import numpy as np +from numpy import typing as npt import pandas as pd from pandas._libs.lib import _NoDefaultDoNotUse @@ -12,29 +15,37 @@ from pandas._typing import ( DtypeBackend, RaiseCoerce, Scalar, - npt, + SequenceNotStr, + np_1darray, ) _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[np.integer]: ... +@overload +def to_numeric( + arg: SequenceNotStr[Any] | npt.NDArray[Any], + errors: RaiseCoerce = "raise", + downcast: _Downcast = None, + dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable", +) -> np_1darray[np.floating]: ... @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..6bd077828 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( @@ -22,4 +22,4 @@ def hash_array( 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/printing.pyi b/pandas-stubs/io/formats/printing.pyi new file mode 100644 index 000000000..ccd0b3dfb --- /dev/null +++ b/pandas-stubs/io/formats/printing.pyi @@ -0,0 +1,6 @@ +from typing import TypeVar + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + +class PrettyDict(dict[_KT, _VT]): ... diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi index 8516c9075..aecf5c8cf 100644 --- a/pandas-stubs/io/formats/style.pyi +++ b/pandas-stubs/io/formats/style.pyi @@ -52,7 +52,7 @@ class _SeriesFunc(Protocol): class _DataFrameFunc(Protocol): def __call__( self, series: DataFrame, /, *args: Any, **kwargs: Any - ) -> npt.NDArray | DataFrame: ... + ) -> npt.NDArray[Any] | DataFrame: ... class _MapCallable(Protocol): def __call__( @@ -238,7 +238,7 @@ class Styler(StylerRenderer): @overload def apply( self, - func: _DataFrameFunc | Callable[[DataFrame], npt.NDArray | DataFrame], + func: _DataFrameFunc | Callable[[DataFrame], npt.NDArray[Any] | DataFrame], axis: None, subset: Subset | None = ..., **kwargs: Any, @@ -295,7 +295,7 @@ class Styler(StylerRenderer): gmap: ( Sequence[float] | Sequence[Sequence[float]] - | npt.NDArray + | npt.NDArray[Any] | DataFrame | Series | None @@ -313,7 +313,7 @@ class Styler(StylerRenderer): gmap: ( Sequence[float] | Sequence[Sequence[float]] - | npt.NDArray + | npt.NDArray[Any] | 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 | npt.NDArray[Any] | DataFrame], float] ) = "mid", vmin: float | None = None, vmax: float | None = None, diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 73bc51bb5..a753fdca5 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -26,6 +26,7 @@ ) import numpy as np +from numpy import typing as npt import pandas as pd from pandas.api.extensions import ( ExtensionArray, @@ -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 @@ -771,7 +773,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[Any]), np_1darray[Any]) def test_types_apply() -> None: @@ -1412,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[Any] | ExtensionArray | pd.Categorical, + ), + np_1darray, + np.integer, ) - valresult_type: type[np.ndarray | ExtensionArray] + valresult_type: type[np_1darray[Any] | ExtensionArray | pd.Categorical] if PD_LTE_23: - valresult_type = np.ndarray + valresult_type = np_1darray[Any] 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[Any] | ExtensionArray | pd.Categorical, + ), valresult_type, + str, ) check( assert_type( pd.Series(list("aabc")).astype("category").values, - np.ndarray | ExtensionArray, + np_1darray[Any] | 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[Any] | ExtensionArray | pd.Categorical, ), - np.ndarray, + np_1darray, + np.datetime64, ) @@ -1572,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), npt.NDArray[Any]), npt.NDArray[Any]) + check(assert_type(s1 @ n1, npt.NDArray[Any]), npt.NDArray[Any]) def test_series_loc_setitem() -> None: diff --git a/tests/test_frame.py b/tests/test_frame.py index 560beff1f..39d73671a 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -54,6 +54,7 @@ TYPE_CHECKING_INVALID_USAGE, check, ensure_clean, + np_1darray, np_2darray, pytest_warns_bounded, ) @@ -626,7 +627,8 @@ 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 | npt.NDArray[Any] | pd.DataFrame | pd.Series, ), pd.DataFrame, ) @@ -1082,7 +1084,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[Any]), np_1darray[Any]) def test_types_apply() -> None: diff --git a/tests/test_pandas.py b/tests/test_pandas.py index af442375b..c3e04ccd6 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -9,14 +9,14 @@ ) import numpy as np -import numpy.typing as npt +from numpy import 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.groupby.grouper import Grouper import pandas.util as pdutil # TODO: pandas-dev/pandas#55023 @@ -558,21 +558,26 @@ def test_read_xml() -> None: def test_unique() -> None: # Taken from the docs check( - assert_type(pd.unique(pd.Series([2, 1, 3, 3])), np.ndarray | ExtensionArray), - np.ndarray, + assert_type( + pd.unique(pd.Series([2, 1, 3, 3])), np_1darray[Any] | ExtensionArray + ), + np_1darray[np.int64], ) check( - assert_type(pd.unique(pd.Series([2] + [1] * 5)), np.ndarray | ExtensionArray), - np.ndarray, + assert_type( + pd.unique(pd.Series([2] + [1] * 5)), np_1darray[Any] | ExtensionArray + ), + np_1darray[np.int64], ) check( assert_type( pd.unique(pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")])), - np.ndarray | ExtensionArray, + np_1darray[Any] | ExtensionArray, ), - np.ndarray, + np_1darray, + np.datetime64, ) check( @@ -585,7 +590,7 @@ def test_unique() -> None: ] ) ), - np.ndarray | ExtensionArray, + np_1darray[Any] | ExtensionArray, ), pd.arrays.DatetimeArray, ) @@ -599,24 +604,27 @@ def test_unique() -> None: ] ) ), - np.ndarray, + np_1darray[np.datetime64] | pd.DatetimeIndex, ), pd.DatetimeIndex, ) - check(assert_type(pd.unique(np.array(list("baabc"))), np.ndarray), np.ndarray) + check( + assert_type(pd.unique(np.array(list("baabc"))), np_1darray[Any]), + np_1darray[Any], + ) check( assert_type( pd.unique(pd.Series(pd.Categorical(list("baabc")))), - np.ndarray | ExtensionArray, + np_1darray[Any] | ExtensionArray, ), pd.Categorical, ) check( assert_type( pd.unique(pd.Series(pd.Categorical(list("baabc"), categories=list("abc")))), - np.ndarray | ExtensionArray, + np_1darray[Any] | ExtensionArray, ), pd.Categorical, ) @@ -627,16 +635,18 @@ def test_unique() -> None: pd.Categorical(list("baabc"), categories=list("abc"), ordered=True) ) ), - np.ndarray | ExtensionArray, + np_1darray[Any] | ExtensionArray, ), pd.Categorical, ) check( - assert_type(pd.unique(pd.Index(["a", "b", "c", "a"])), np.ndarray), + assert_type( + pd.unique(pd.Index(["a", "b", "c", "a"])), np_1darray[Any] | pd.Index + ), np.ndarray if PD_LTE_23 else pd.Index, ) check( - assert_type(pd.unique(pd.RangeIndex(0, 10)), np.ndarray), + assert_type(pd.unique(pd.RangeIndex(0, 10)), np_1darray[Any] | pd.Index), np.ndarray if PD_LTE_23 else pd.Index, ) check( @@ -653,7 +663,7 @@ def test_unique() -> None: check( assert_type( pd.unique(pd.timedelta_range(start="1 day", periods=4)), - np.ndarray, + np_1darray[Any] | pd.Index, ), np.ndarray if PD_LTE_23 else pd.Index, ) @@ -692,13 +702,15 @@ 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) @@ -723,7 +735,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, + npt.NDArray[Any] | Scalar | pd.DataFrame | pd.Series | None, ), pd.DataFrame, ) @@ -741,50 +753,46 @@ 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[np.integer]), + 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[np.floating]), + 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[np.floating]), + np_1darray, + np.floating, ) check( assert_type( pd.to_numeric(np.array([1.0, 2.0, "3.0"], dtype=object)), - npt.NDArray, + np_1darray[np.floating], ), - 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[np.floating] ), - 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[np.floating]), + 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[np.integer] + ), + np_1darray, + np.integer, ) @@ -931,34 +939,34 @@ 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[Any]), np_1darray[Any]) 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[Any]), np_1darray[Any]) 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[Any]), np_1darray[Any]) 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[Any]), np_1darray[Any]) def test_index_unqiue() -> None: @@ -976,18 +984,31 @@ 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[np.datetime64] | pd.DatetimeIndex), + np_1darray if PD_LTE_23 else pd.DatetimeIndex, + np.datetime64, + ) + check( + assert_type(pd.unique(i), np_1darray[Any] | pd.Index), + np_1darray[Any] if PD_LTE_23 else pd.Index, ) - check(assert_type(pd.unique(i), np.ndarray), np.ndarray if PD_LTE_23 else 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[Any] | pd.Index), + np_1darray[Any] if PD_LTE_23 else pd.Index, + ) + check( + assert_type(pd.unique(tdi), np_1darray[Any] | pd.Index), + np_1darray[Any] if PD_LTE_23 else pd.Index, + ) + check( + assert_type(pd.unique(mi), np_1darray[Any] | pd.Index), + np_1darray[Any] if PD_LTE_23 else pd.Index, ) - 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, ) @@ -997,7 +1018,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) @@ -1006,11 +1030,14 @@ 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[np.double]), np_1darray[np.double]) + check( + assert_type(e0, np_1darray[np.intp]), + np_1darray[np.intp], + ) + check(assert_type(e1, np_1darray[np.double]), np_1darray[np.double]) check(assert_type(f0, pd.Categorical), pd.Categorical) - check(assert_type(f1, npt.NDArray), np.ndarray) + check(assert_type(f1, np_1darray[np.double]), np_1darray[np.double]) 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") @@ -1045,11 +1072,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[np.double]), np_1darray[np.double]) check(assert_type(k0, pd.Series), pd.Series) - check(assert_type(k1, npt.NDArray), np.ndarray) + check(assert_type(k1, np_1darray[np.double]), np_1darray[np.double]) check(assert_type(l0, pd.Series), pd.Series) - check(assert_type(l1, npt.NDArray), np.ndarray) + check(assert_type(l1, np_1darray[np.double]), np_1darray[np.double]) check(assert_type(m0, pd.Series), pd.Series) check(assert_type(m1, pd.IntervalIndex), pd.IntervalIndex) @@ -1122,10 +1149,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[np.double]), np_1darray[np.double]) + check(assert_type(b1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(c1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(d1, np_1darray[np.double]), np_1darray[np.double]) 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) @@ -1136,17 +1163,23 @@ 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[np.intp | np.double]), + np_1darray[np.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) + check( + assert_type(i0, np_1darray[np.intp | np.double]), + np_1darray[np.double], # because of nans + ) + check(assert_type(j0, np_1darray[np.intp | np.double]), np_1darray[np.double]) + + check(assert_type(e1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(f1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(g1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(h1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(i1, np_1darray[np.double]), np_1darray[np.double]) + check(assert_type(j1, np_1darray[np.double]), np_1darray[np.double]) @pytest.mark.xfail(