From e20a0606a1475e369f8cbb9d2a47214f67ae26d7 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Sat, 1 Nov 2025 00:25:55 +0100 Subject: [PATCH] wip --- pandas-stubs/_libs/missing.pyi | 131 +++++---- pandas-stubs/_typing.pyi | 82 +++--- pandas-stubs/core/algorithms.pyi | 4 +- pandas-stubs/core/arrays/arrow/array.pyi | 6 + pandas-stubs/core/arrays/numpy_.pyi | 3 + pandas-stubs/core/base.pyi | 3 +- pandas-stubs/core/indexes/accessors.pyi | 17 +- pandas-stubs/core/indexes/base.pyi | 9 +- pandas-stubs/core/reshape/tile.pyi | 18 +- pandas-stubs/core/series.pyi | 327 +++++++++++++++-------- tests/indexes/test_indexes.py | 8 +- tests/series/test_properties.py | 12 +- tests/series/test_series.py | 47 ++-- tests/test_natype.py | 58 ++-- tests/test_pandas.py | 22 +- 15 files changed, 460 insertions(+), 287 deletions(-) create mode 100644 pandas-stubs/core/arrays/arrow/array.pyi diff --git a/pandas-stubs/_libs/missing.pyi b/pandas-stubs/_libs/missing.pyi index f0bf682e3..59ef27bba 100644 --- a/pandas-stubs/_libs/missing.pyi +++ b/pandas-stubs/_libs/missing.pyi @@ -5,13 +5,16 @@ from typing import ( overload, ) -from pandas import ( - Index, - Series, -) from pandas.core.arrays.boolean import BooleanArray +from pandas.core.indexes.base import Index +from pandas.core.series import Series from typing_extensions import Self +from pandas._typing import ( + S1, + A1N_co, +) + class NAType: def __new__(cls, *args: Any, **kwargs: Any) -> Self: ... def __format__(self, format_spec: str) -> str: ... @@ -19,48 +22,60 @@ class NAType: def __reduce__(self) -> str: ... @overload def __add__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __add__(self, other: Series, /) -> Series: ... @overload def __add__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __add__(self, other: object, /) -> NAType: ... @overload def __radd__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __radd__(self, other: Series, /) -> Series: ... @overload def __radd__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __radd__(self, other: object, /) -> NAType: ... @overload def __sub__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __sub__(self, other: Series, /) -> Series: ... @overload def __sub__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __sub__(self, other: object, /) -> NAType: ... @overload def __rsub__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rsub__(self, other: Series, /) -> Series: ... @overload def __rsub__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __rsub__(self, other: object, /) -> NAType: ... @overload def __mul__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __mul__(self, other: Series, /) -> Series: ... @overload def __mul__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __mul__(self, other: object, /) -> NAType: ... @overload def __rmul__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rmul__(self, other: Series, /) -> Series: ... @overload def __rmul__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload @@ -69,48 +84,60 @@ class NAType: def __rmatmul__(self, other: object, /) -> NAType: ... @overload def __truediv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __truediv__(self, other: Series, /) -> Series: ... @overload def __truediv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __truediv__(self, other: object, /) -> NAType: ... @overload def __rtruediv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rtruediv__(self, other: Series, /) -> Series: ... @overload def __rtruediv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __rtruediv__(self, other: object, /) -> NAType: ... @overload def __floordiv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __floordiv__(self, other: Series, /) -> Series: ... @overload def __floordiv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __floordiv__(self, other: object, /) -> NAType: ... @overload def __rfloordiv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rfloordiv__(self, other: Series, /) -> Series: ... @overload def __rfloordiv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __rfloordiv__(self, other: object, /) -> NAType: ... @overload def __mod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __mod__(self, other: Series, /) -> Series: ... @overload def __mod__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __mod__(self, other: object, /) -> NAType: ... @overload def __rmod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rmod__(self, other: Series, /) -> Series: ... @overload def __rmod__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload @@ -133,8 +160,8 @@ class NAType: def __rdivmod__(self, other: object, /) -> tuple[NAType, NAType]: ... @overload # type: ignore[override] def __eq__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __eq__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload @@ -143,8 +170,8 @@ class NAType: ) -> NAType: ... @overload # type: ignore[override] def __ne__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __ne__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload @@ -153,32 +180,32 @@ class NAType: ) -> NAType: ... @overload def __le__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __le__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload def __le__(self, other: object, /) -> NAType: ... @overload def __lt__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __lt__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload def __lt__(self, other: object, /) -> NAType: ... @overload def __gt__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __gt__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload def __gt__(self, other: object, /) -> NAType: ... @overload def __ge__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series[bool]: ... + self, other: Series[Any, Any], / + ) -> Series[bool, BooleanArray]: ... @overload def __ge__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] @overload @@ -189,16 +216,20 @@ class NAType: def __invert__(self) -> NAType: ... @overload def __pow__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __pow__(self, other: Series, /) -> Series: ... @overload def __pow__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __pow__(self, other: object, /) -> NAType: ... @overload def __rpow__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rpow__(self, other: Series, /) -> Series: ... @overload def __rpow__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload @@ -221,16 +252,20 @@ class NAType: def __ror__(self, other: bool | NAType, /) -> NAType: ... @overload def __xor__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __xor__(self, other: Series, /) -> Series: ... @overload def __xor__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload def __xor__(self, other: object, /) -> NAType: ... @overload def __rxor__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] - self, other: Series, / - ) -> Series: ... + self, other: Series[S1, A1N_co], / + ) -> Series[S1, A1N_co]: ... + @overload + def __rxor__(self, other: Series, /) -> Series: ... @overload def __rxor__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] @overload diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 5f0277185..b830ab8f6 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -33,6 +33,7 @@ from pandas.core.arrays import ( ExtensionArray, IntegerArray, ) +from pandas.core.arrays.numpy_ import NumpyExtensionArray from pandas.core.frame import DataFrame from pandas.core.generic import NDFrame from pandas.core.groupby.grouper import Grouper @@ -222,34 +223,24 @@ Dtype: TypeAlias = ExtensionDtype | NpDtype # object0 # m # datetime64 - +# Builtin bool type and its string alias +PyBooleanDtypeArg: TypeAlias = type[bool] | Literal["bool"] +# Pandas nullable boolean type and its string alias +PdBooleanDtypeArg: TypeAlias = pd.BooleanDtype | Literal["boolean"] +# Numpy bool type +# https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bool_ +NpBooleanDtypeArg: TypeAlias = type[np.bool_] | Literal["?", "b1", "bool_"] +# PyArrow boolean type and its string alias +PaBooleanDtypeArg: TypeAlias = Literal["bool[pyarrow]", "boolean[pyarrow]"] BooleanDtypeArg: TypeAlias = ( - # Builtin bool type and its string alias - type[bool] # noqa: PYI030 - | Literal["bool"] - # Pandas nullable boolean type and its string alias - | pd.BooleanDtype - | Literal["boolean"] - # Numpy bool type - # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bool_ - | type[np.bool_] - | Literal["?", "b1", "bool_"] - # PyArrow boolean type and its string alias - | Literal["bool[pyarrow]", "boolean[pyarrow]"] + PyBooleanDtypeArg | PdBooleanDtypeArg | NpBooleanDtypeArg | PaBooleanDtypeArg ) -IntDtypeArg: TypeAlias = ( - # Builtin integer type and its string alias - type[int] # noqa: PYI030 - | Literal["int"] - # Pandas nullable integer types and their string aliases - | pd.Int8Dtype - | pd.Int16Dtype - | pd.Int32Dtype - | pd.Int64Dtype - | Literal["Int8", "Int16", "Int32", "Int64"] - # Numpy signed integer types and their string aliases +# Builtin integer type and its string alias +PyIntDtypeArg: TypeAlias = type[int] | Literal["int"] +# Numpy signed integer types and their string aliases +NpIntDtypeArg: TypeAlias = ( # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.byte - | type[np.byte] + type[np.byte] # noqa: PYI030 | Literal["b", "i1", "int8", "byte"] # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.short | type[np.short] @@ -266,19 +257,32 @@ IntDtypeArg: TypeAlias = ( # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.intp | type[np.intp] # signed pointer (=`intptr_t`, platform dependent) | Literal["p", "intp"] - # PyArrow integer types and their string aliases - | Literal["int8[pyarrow]", "int16[pyarrow]", "int32[pyarrow]", "int64[pyarrow]"] ) -UIntDtypeArg: TypeAlias = ( - # Pandas nullable unsigned integer types and their string aliases - pd.UInt8Dtype # noqa: PYI030 +# Pandas nullable integer types and their string aliases +PdIntDtypeArg: TypeAlias = ( + pd.Int8Dtype + | pd.Int16Dtype + | pd.Int32Dtype + | pd.Int64Dtype + | Literal["Int8", "Int16", "Int32", "Int64"] +) +# PyArrow integer types and their string aliases +PaIntDtypeArg: TypeAlias = Literal[ + "int8[pyarrow]", "int16[pyarrow]", "int32[pyarrow]", "int64[pyarrow]" +] +IntDtypeArg: TypeAlias = PyIntDtypeArg | NpIntDtypeArg | PdIntDtypeArg | PaIntDtypeArg +# Pandas nullable unsigned integer types and their string aliases +PdUIntDtypeArg: TypeAlias = ( + pd.UInt8Dtype | pd.UInt16Dtype | pd.UInt32Dtype | pd.UInt64Dtype | Literal["UInt8", "UInt16", "UInt32", "UInt64"] - # Numpy unsigned integer types and their string aliases +) +# Numpy unsigned integer types and their string aliases +NpUIntDtypeArg: TypeAlias = ( # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.ubyte - | type[np.ubyte] + type[np.ubyte] # noqa: PYI030 | Literal["B", "u1", "uint8", "ubyte"] # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.ushort | type[np.ushort] @@ -295,9 +299,12 @@ UIntDtypeArg: TypeAlias = ( # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.uintp | type[np.uintp] # unsigned pointer (=`uintptr_t`, platform dependent) | Literal["P", "uintp"] - # PyArrow unsigned integer types and their string aliases - | Literal["uint8[pyarrow]", "uint16[pyarrow]", "uint32[pyarrow]", "uint64[pyarrow]"] ) +# PyArrow unsigned integer types and their string aliases +PaUIntDtypeArg: TypeAlias = Literal[ + "uint8[pyarrow]", "uint16[pyarrow]", "uint32[pyarrow]", "uint64[pyarrow]" +] +UIntDtypeArg: TypeAlias = PdUIntDtypeArg | NpUIntDtypeArg | PaUIntDtypeArg FloatDtypeArg: TypeAlias = ( # Builtin float type and its string alias type[float] # noqa: PYI030 @@ -938,6 +945,13 @@ C2 = TypeVar( BaseOffset, ) +A1N_co = TypeVar( + "A1N_co", bound=ExtensionArray, default=NumpyExtensionArray, covariant=True +) +A1_co = TypeVar( + "A1_co", bound=np.ndarray | ExtensionArray, default=np.ndarray, covariant=True +) + IndexingInt: TypeAlias = ( int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8 ) diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 5174218c6..5f3d35517 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -34,11 +34,11 @@ def unique(values: CategoricalIndex) -> CategoricalIndex: ... # type: ignore[ov @overload def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ... @overload -def unique(values: Index) -> np.ndarray: ... +def unique(values: Index[Any, Any]) -> np.ndarray: ... @overload def unique(values: Categorical) -> Categorical: ... @overload -def unique(values: Series) -> np.ndarray | ExtensionArray: ... +def unique(values: Series[Any, Any]) -> np.ndarray | ExtensionArray: ... @overload def unique(values: np.ndarray) -> np.ndarray: ... @overload diff --git a/pandas-stubs/core/arrays/arrow/array.pyi b/pandas-stubs/core/arrays/arrow/array.pyi new file mode 100644 index 000000000..f913d3c4f --- /dev/null +++ b/pandas-stubs/core/arrays/arrow/array.pyi @@ -0,0 +1,6 @@ +from pandas.core.arraylike import OpsMixin + +from pandas.core.arrays.base import ExtensionArray + + +class ArrowExtensionArray(OpsMixin, ExtensionArray): ... diff --git a/pandas-stubs/core/arrays/numpy_.pyi b/pandas-stubs/core/arrays/numpy_.pyi index 760d82e84..9d10b8a5a 100644 --- a/pandas-stubs/core/arrays/numpy_.pyi +++ b/pandas-stubs/core/arrays/numpy_.pyi @@ -1,5 +1,6 @@ import numpy as np from numpy.lib.mixins import NDArrayOperatorsMixin +from pandas.core.arraylike import OpsMixin from pandas.core.arrays.base import ( ExtensionArray, ExtensionOpsMixin, @@ -7,6 +8,8 @@ from pandas.core.arrays.base import ( from pandas.core.dtypes.dtypes import ExtensionDtype +class NumpyExtensionArray(OpsMixin, ExtensionArray): ... + class PandasDtype(ExtensionDtype): @property def numpy_dtype(self) -> np.dtype: ... diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index ea3865add..96ed0c28a 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -29,6 +29,7 @@ from pandas._libs.tslibs.timedeltas import Timedelta from pandas._typing import ( S1, S2, + A1_co, AxisIndex, DropKeep, DTypeLike, @@ -59,7 +60,7 @@ class SelectionMixin(Generic[NDFrameT]): def __getitem__(self, key): ... def aggregate(self, func, *args: Any, **kwargs: Any): ... -class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]): +class IndexOpsMixin(OpsMixin, Generic[S1, A1_co, GenericT_co]): __array_priority__: int = ... @property def T(self) -> Self: ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index b87872d78..892f7f9bd 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -5,6 +5,7 @@ from datetime import ( tzinfo as _tzinfo, ) from typing import ( + Any, Generic, Literal, TypeVar, @@ -14,10 +15,10 @@ from typing import ( import numpy as np from pandas.core.accessor import PandasDelegate -from pandas.core.arrays.base import ExtensionArray from pandas.core.arrays.categorical import Categorical from pandas.core.arrays.datetimes import DatetimeArray from pandas.core.arrays.interval import IntervalArray +from pandas.core.arrays.numpy_ import NumpyExtensionArray from pandas.core.arrays.period import PeriodArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.base import ( @@ -49,8 +50,6 @@ from pandas._typing import ( np_ndarray_bool, ) -from pandas.core.dtypes.dtypes import CategoricalDtype - class Properties(PandasDelegate, NoNewAttributesMixin): ... _DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", bound=Series[int] | Index[int]) @@ -462,11 +461,9 @@ class DtDescriptor: class ArrayDescriptor: @overload def __get__( - self, instance: IndexOpsMixin[Never], owner: type[IndexOpsMixin] - ) -> ExtensionArray: ... - @overload - def __get__( - self, instance: IndexOpsMixin[CategoricalDtype], owner: type[IndexOpsMixin] + self, + instance: IndexOpsMixin[Any, Categorical], + owner: type[IndexOpsMixin[Any, Categorical]], ) -> Categorical: ... @overload def __get__( @@ -483,5 +480,5 @@ class ArrayDescriptor: # should be NumpyExtensionArray @overload def __get__( - self, instance: IndexOpsMixin, owner: type[IndexOpsMixin] - ) -> ExtensionArray: ... + self, instance: IndexOpsMixin[Any, np.ndarray], owner: type[IndexOpsMixin] + ) -> NumpyExtensionArray: ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 21aa39db3..2931efbb1 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -67,6 +67,7 @@ from pandas._typing import ( S2, S2_NSDT, T_COMPLEX, + A1_co, AnyAll, AnyArrayLike, AnyArrayLikeInt, @@ -109,7 +110,7 @@ from pandas._typing import ( class InvalidIndexError(Exception): ... -class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): +class Index(IndexOpsMixin[S1, A1_co], ElementOpsMixin[S1]): __hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride] # overloads with additional dtypes @overload @@ -473,7 +474,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): def values(self) -> np_1darray: ... def memory_usage(self, deep: bool = False): ... @overload - def where( + def where( # pyright: ignore[reportOverlappingOverload] self, cond: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool], other: S1 | Series[S1] | Self, @@ -541,7 +542,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]): @overload def insert(self, loc: int, item: S1) -> Self: ... @overload - def insert(self, loc: int, item: object) -> Index: ... + def insert( + self: Index[Any, A1_co], loc: int, item: object + ) -> Index[Any, A1_co]: ... def drop(self, labels, errors: IgnoreRaise = "raise") -> Self: ... @property def shape(self) -> tuple[int, ...]: ... diff --git a/pandas-stubs/core/reshape/tile.pyi b/pandas-stubs/core/reshape/tile.pyi index 0eed14683..cac2eeadc 100644 --- a/pandas-stubs/core/reshape/tile.pyi +++ b/pandas-stubs/core/reshape/tile.pyi @@ -7,7 +7,6 @@ from typing import ( import numpy as np from pandas import ( Categorical, - CategoricalDtype, DatetimeIndex, Index, Interval, @@ -16,6 +15,7 @@ from pandas import ( ) from pandas.core.series import Series +from pandas._libs.interval import _OrderableT from pandas._typing import ( IntervalT, Label, @@ -171,19 +171,11 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> Series[CategoricalDtype]: ... +) -> Series[Interval[Timestamp], Categorical]: ... @overload def cut( - x: Series, - bins: ( - int - | Series - | Index[int] - | Index[float] - | Sequence[int] - | Sequence[float] - | IntervalIndex - ), + x: Series[_OrderableT], + bins: int | Series | Index[int] | Index[float] | Sequence[float] | IntervalIndex, right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., retbins: Literal[False] = False, @@ -191,7 +183,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> Series: ... +) -> Series[Interval[_OrderableT], Categorical]: ... @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 32b6bc8a5..c7ec997da 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,9 +61,14 @@ from pandas.core.api import ( Int32Dtype as Int32Dtype, Int64Dtype as Int64Dtype, ) -from pandas.core.arrays.boolean import BooleanDtype -from pandas.core.arrays.categorical import CategoricalAccessor +from pandas.core.arrays.arrow.array import ArrowExtensionArray +from pandas.core.arrays.boolean import BooleanArray +from pandas.core.arrays.categorical import ( + Categorical, + CategoricalAccessor, +) from pandas.core.arrays.datetimes import DatetimeArray +from pandas.core.arrays.integer import IntegerArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.base import ( ElementOpsMixin, @@ -122,6 +128,8 @@ from pandas._typing import ( S2, S2_NSDT, T_COMPLEX, + A1_co, + A1N_co, AggFuncTypeBase, AggFuncTypeDictFrame, AggFuncTypeSeriesToFrame, @@ -170,9 +178,20 @@ from pandas._typing import ( ListLikeU, MaskType, NaPosition, + NpBooleanDtypeArg, + NpIntDtypeArg, + NpUIntDtypeArg, NsmallestNlargestKeep, ObjectDtypeArg, + PaBooleanDtypeArg, + PaIntDtypeArg, + PaUIntDtypeArg, + PdBooleanDtypeArg, + PdIntDtypeArg, + PdUIntDtypeArg, PeriodFrequency, + PyBooleanDtypeArg, + PyIntDtypeArg, QuantileInterpolation, RandomState, ReindexMethod, @@ -210,12 +229,10 @@ from pandas._typing import ( np_ndarray_float, np_ndarray_str, np_ndarray_td, - npt, num, ) from pandas.core.dtypes.base import ExtensionDtype -from pandas.core.dtypes.dtypes import CategoricalDtype from pandas.plotting import PlotAccessor @@ -245,7 +262,7 @@ class _iLocIndexerSeries(_iLocIndexer, Generic[S1]): @overload def __getitem__( self, idx: Index | Series | slice | np_ndarray_anyint - ) -> Series[S1]: ... + ) -> Self: ... # set item # Keep in sync with `Series.__setitem__` @@ -283,7 +300,7 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): ), # _IndexSliceTuple is when having a tuple that includes a slice. Could just # be s.loc[1, :], or s.loc[pd.IndexSlice[1, :]] - ) -> Series[S1]: ... + ) -> Self: ... # Keep in sync with `Series.__setitem__` @overload @@ -310,20 +327,11 @@ _DataLikeS1: TypeAlias = ( ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | IndexOpsMixin[S1] ) -class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): +class Series(IndexOpsMixin[S1, A1_co], ElementOpsMixin[S1], NDFrame): # Define __index__ because mypy thinks Series follows protocol `SupportsIndex` https://github.com/pandas-dev/pandas-stubs/pull/1332#discussion_r2285648790 __index__: ClassVar[None] __hash__: ClassVar[None] # pyright: ignore[reportIncompatibleMethodOverride] - @overload - def __new__( - cls, - data: npt.NDArray[np.float64], - index: AxesData | None = ..., - dtype: Dtype = ..., - name: Hashable = ..., - copy: bool = ..., - ) -> Series[float]: ... @overload def __new__( cls, @@ -378,16 +386,6 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = ..., ) -> Series[Timestamp]: ... @overload - def __new__( - cls, - data: _DataLike, - index: AxesData | None = ..., - *, - dtype: CategoryDtypeArg, - name: Hashable = ..., - copy: bool = ..., - ) -> Series[CategoricalDtype]: ... - @overload def __new__( cls, data: PeriodIndex | Sequence[Period], @@ -435,7 +433,33 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = ..., ) -> Series[Interval[_OrderableT]]: ... @overload - def __new__( # type: ignore[overload-overlap] + def __new__( # pyright: ignore[reportOverlappingOverload] + cls, + data: Categorical, + index: AxesData | None = ..., + dtype: Dtype = ..., + name: Hashable = ..., + copy: bool = ..., + ) -> Series[Any, Categorical]: ... + @overload + def __new__( # pyright: ignore[reportOverlappingOverload] + cls, + data: ( + S1 + | Sequence[S1] + | dict[HashableT1, S1] + | KeysView[S1] + | ValuesView[S1] + | IndexOpsMixin[S1] + ), + index: AxesData | None = ..., + *, + dtype: CategoryDtypeArg, + name: Hashable = ..., + copy: bool = ..., + ) -> Series[S1, Categorical]: ... + @overload + def __new__( # pyright: ignore[reportOverlappingOverload] cls, data: Scalar | _DataLike | dict[HashableT1, Any] | None, index: AxesData | None = ..., @@ -445,36 +469,123 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: bool = ..., ) -> Self: ... @overload - def __new__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + def __new__( cls, - data: Sequence[bool], + data: ( + bool + | Sequence[bool] + | dict[HashableT1, bool] + | KeysView[bool] + | ValuesView[bool] + | npt.NDArray[np.bool] + | IndexOpsMixin[bool] + ), index: AxesData | None = ..., - dtype: Dtype = ..., + dtype: PyBooleanDtypeArg | NpBooleanDtypeArg = ..., name: Hashable = ..., copy: bool = ..., ) -> Series[bool]: ... @overload + def __new__( + cls, + data: ( + bool + | Sequence[bool] + | dict[HashableT1, bool] + | KeysView[bool] + | ValuesView[bool] + | npt.NDArray[np.bool] + | IndexOpsMixin[bool] + ), + index: AxesData | None = ..., + *, + dtype: PdBooleanDtypeArg, + name: Hashable = ..., + copy: bool = ..., + ) -> Series[bool, BooleanArray]: ... + @overload + def __new__( + cls, + data: ( + bool + | Sequence[bool] + | dict[HashableT1, bool] + | KeysView[bool] + | ValuesView[bool] + | npt.NDArray[np.bool] + | IndexOpsMixin[bool] + ), + index: AxesData | None = ..., + *, + dtype: PaBooleanDtypeArg, + name: Hashable = ..., + copy: bool = ..., + ) -> Series[bool, ArrowExtensionArray]: ... + @overload def __new__( # type: ignore[overload-overlap] cls, - data: Sequence[int], + data: ( + int + | Sequence[int] + | dict[HashableT1, int] + | KeysView[int] + | ValuesView[int] + | npt.NDArray[np.integer] + | IndexOpsMixin[int] + ), index: AxesData | None = ..., - dtype: Dtype = ..., + dtype: PyIntDtypeArg | NpIntDtypeArg | NpUIntDtypeArg = ..., name: Hashable = ..., copy: bool = ..., ) -> Series[int]: ... @overload - def __new__( + def __new__( # type: ignore[overload-overlap] cls, - data: Sequence[float], + data: ( + int + | Sequence[int] + | dict[HashableT1, int] + | KeysView[int] + | ValuesView[int] + | npt.NDArray[np.integer] + | IndexOpsMixin[int] + ), index: AxesData | None = ..., - dtype: Dtype = ..., + *, + dtype: PdIntDtypeArg | PdUIntDtypeArg, name: Hashable = ..., copy: bool = ..., - ) -> Series[float]: ... + ) -> Series[int, IntegerArray]: ... @overload - def __new__( # type: ignore[overload-cannot-match] # pyright: ignore[reportOverlappingOverload] + def __new__( # type: ignore[overload-overlap] cls, - data: Sequence[int | float], + data: ( + int + | Sequence[int] + | dict[HashableT1, int] + | KeysView[int] + | ValuesView[int] + | npt.NDArray[np.integer] + | IndexOpsMixin[int] + ), + index: AxesData | None = ..., + *, + dtype: PaIntDtypeArg | PaUIntDtypeArg, + name: Hashable = ..., + copy: bool = ..., + ) -> Series[int, ArrowExtensionArray]: ... + @overload + def __new__( + cls, + data: ( + float + | Sequence[float] + | dict[HashableT1, float] + | KeysView[float] + | ValuesView[float] + | npt.NDArray[np.floating] + | IndexOpsMixin[float] + ), index: AxesData | None = ..., dtype: Dtype = ..., name: Hashable = ..., @@ -522,7 +633,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def values(self) -> ArrayLike: ... def ravel(self, order: _str = ...) -> np.ndarray: ... def __len__(self) -> int: ... - def view(self, dtype: Dtype | None = None) -> Series[S1]: ... + def view(self, dtype: Dtype | None = None) -> Self: ... @final def __array_ufunc__( self, ufunc: Callable, method: _str, *inputs: Any, **kwargs: Any @@ -600,9 +711,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): def get(self, key: Hashable, default: S1) -> S1: ... @overload def get(self, key: Hashable, default: _T) -> S1 | _T: ... - def repeat( - self, repeats: int | list[int], axis: AxisIndex | None = 0 - ) -> Series[S1]: ... + def repeat(self, repeats: int | list[int], axis: AxisIndex | None = 0) -> Self: ... @property def index(self) -> Index: ... @index.setter @@ -628,7 +737,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): name: Level = ..., inplace: Literal[False] = False, allow_duplicates: bool = ..., - ) -> Series[S1]: ... + ) -> Self: ... @overload def reset_index( self, @@ -858,7 +967,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): dropna: _bool = ..., ) -> SeriesGroupBy[S1, Any]: ... def count(self) -> int: ... - def mode(self, dropna: bool = True) -> Series[S1]: ... + def mode(self, dropna: bool = True) -> Self: ... @overload def unique(self: Series[Never]) -> np.ndarray: ... # type: ignore[overload-overlap] @overload @@ -882,7 +991,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): keep: DropKeep = ..., inplace: Literal[False] = False, ignore_index: _bool = ..., - ) -> Series[S1]: ... + ) -> Self: ... def duplicated(self, keep: DropKeep = "first") -> Series[_bool]: ... def idxmax( self, @@ -898,7 +1007,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): *args: Any, **kwargs: Any, ) -> int | _str: ... - def round(self, decimals: int = 0, *args: Any, **kwargs: Any) -> Series[S1]: ... + def round(self, decimals: int = 0, *args: Any, **kwargs: Any) -> Self: ... @overload def quantile( self, @@ -910,7 +1019,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): self, q: ListLike, interpolation: QuantileInterpolation = ..., - ) -> Series[S1]: ... + ) -> Self: ... def corr( self, other: Series[S1], @@ -921,16 +1030,18 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): self, other: Series[S1], min_periods: int | None = None, ddof: int = 1 ) -> float: ... @overload + def diff( + self: Series[bool, A1N_co], periods: int = ... + ) -> Series[bool, A1N_co]: ... + @overload + def diff(self: Series[int, A1N_co], periods: int = ...) -> Series[int, A1N_co]: ... + @overload def diff( # type: ignore[overload-overlap] self: Series[Never] | Series[int], periods: int = ... ) -> Series[float]: ... @overload def diff(self: Series[_bool], periods: int = ...) -> Series: ... @overload - def diff( - self: Series[BooleanDtype], periods: int = ... - ) -> Series[BooleanDtype]: ... - @overload def diff(self: Series[Interval], periods: int = ...) -> Never: ... @overload def diff( @@ -940,7 +1051,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @overload def dot(self, other: Series[S1]) -> Scalar: ... @overload - def dot(self, other: DataFrame) -> Series[S1]: ... + def dot(self, other: DataFrame) -> Self: ... @overload def dot( self, other: ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | Index[S1] @@ -991,8 +1102,8 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ) -> DataFrame: ... def combine( self, other: Series[S1], func: Callable, fill_value: Scalar | None = ... - ) -> Series[S1]: ... - def combine_first(self, other: Series[S1]) -> Series[S1]: ... + ) -> Self: ... + def combine_first(self, other: Series[S1]) -> Self: ... def update(self, other: Series[S1] | Sequence[S1] | Mapping[int, S1]) -> None: ... @overload def sort_values( @@ -1017,7 +1128,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ignore_index: _bool = ..., inplace: Literal[False] = False, key: ValueKeyFunc = ..., - ) -> Series[S1]: ... + ) -> Self: ... @overload def sort_index( self, @@ -1045,7 +1156,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): ignore_index: _bool = ..., inplace: Literal[False] = False, key: IndexKeyFunc = ..., - ) -> Series[S1]: ... + ) -> Self: ... def argsort( self, axis: AxisIndex = 0, @@ -1053,17 +1164,11 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): order: None = None, stable: None = None, ) -> Series[int]: ... - def nlargest( - self, n: int = 5, keep: NsmallestNlargestKeep = "first" - ) -> Series[S1]: ... - def nsmallest( - self, n: int = 5, keep: NsmallestNlargestKeep = "first" - ) -> Series[S1]: ... - def swaplevel( - self, i: Level = -2, j: Level = -1, copy: _bool = True - ) -> Series[S1]: ... - def reorder_levels(self, order: list) -> Series[S1]: ... - def explode(self, ignore_index: _bool = ...) -> Series[S1]: ... + def nlargest(self, n: int = 5, keep: NsmallestNlargestKeep = "first") -> Self: ... + def nsmallest(self, n: int = 5, keep: NsmallestNlargestKeep = "first") -> Self: ... + def swaplevel(self, i: Level = -2, j: Level = -1, copy: _bool = True) -> Self: ... + def reorder_levels(self, order: list) -> Self: ... + def explode(self, ignore_index: _bool = ...) -> Self: ... def unstack( self, level: IndexLabel = -1, @@ -1120,7 +1225,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): axis: AxisIndex = ..., *args: Any, **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... @overload def transform( self, @@ -1235,7 +1340,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): axis: AxisIndex = ..., limit: int | None = ..., inplace: Literal[False] = False, - ) -> Series[S1]: ... + ) -> Self: ... @overload def replace( self, @@ -1253,7 +1358,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): *, regex: ReplaceValue = ..., inplace: Literal[False] = False, - ) -> Series[S1]: ... + ) -> Self: ... def shift( self, periods: int | Sequence[int] = ..., @@ -1297,13 +1402,13 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): inplace: Literal[False] = False, how: AnyAll | None = ..., ignore_index: _bool = ..., - ) -> Series[S1]: ... + ) -> Self: ... def to_timestamp( self, freq: PeriodFrequency | None = None, how: ToTimestampHow = "start", copy: _bool = True, - ) -> Series[S1]: ... + ) -> Self: ... def to_period( self, freq: PeriodFrequency | None = None, copy: _bool = True ) -> DataFrame: ... @@ -1342,18 +1447,18 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): @final def swapaxes( self, axis1: AxisIndex, axis2: AxisIndex, copy: _bool = ... - ) -> Series[S1]: ... + ) -> Self: ... @final def droplevel(self, level: Level | list[Level], axis: AxisIndex = 0) -> Self: ... def pop(self, item: Hashable) -> S1: ... @final def squeeze(self, axis: None = None) -> Series[S1] | Scalar: ... @final - def __abs__(self) -> Series[S1]: ... + def __abs__(self) -> Self: ... @final - def add_prefix(self, prefix: _str, axis: AxisIndex | None = None) -> Series[S1]: ... + def add_prefix(self, prefix: _str, axis: AxisIndex | None = None) -> Self: ... @final - def add_suffix(self, suffix: _str, axis: AxisIndex | None = None) -> Series[S1]: ... + def add_suffix(self, suffix: _str, axis: AxisIndex | None = None) -> Self: ... def reindex( self, index: Axes | None = None, @@ -1363,18 +1468,18 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): fill_value: Scalar | None = None, limit: int | None = None, tolerance: float | Timedelta | None = None, - ) -> Series[S1]: ... + ) -> Self: ... def filter( self, items: ListLike | None = None, like: _str | None = None, regex: _str | None = None, axis: AxisIndex | None = None, - ) -> Series[S1]: ... + ) -> Self: ... @final - def head(self, n: int = 5) -> Series[S1]: ... + def head(self, n: int = 5) -> Self: ... @final - def tail(self, n: int = 5) -> Series[S1]: ... + def tail(self, n: int = 5) -> Self: ... @final def sample( self, @@ -1385,7 +1490,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): random_state: RandomState | None = None, axis: AxisIndex | None = None, ignore_index: _bool = False, - ) -> Series[S1]: ... + ) -> Self: ... @overload def astype( self, @@ -1448,7 +1553,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): dtype: CategoryDtypeArg, copy: _bool = ..., errors: IgnoreRaise = ..., - ) -> Series[CategoricalDtype]: ... + ) -> Series[S1, Categorical]: ... @overload def astype( self, @@ -1457,9 +1562,9 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): errors: IgnoreRaise = ..., ) -> Series: ... @final - def copy(self, deep: _bool = True) -> Series[S1]: ... + def copy(self, deep: _bool = True) -> Self: ... @final - def infer_objects(self, copy: _bool = True) -> Series[S1]: ... + def infer_objects(self, copy: _bool = True) -> Self: ... @overload def ffill( self, @@ -1477,7 +1582,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): inplace: Literal[False] = False, limit: int | None = ..., limit_area: Literal["inside", "outside"] | None = ..., - ) -> Series[S1]: ... + ) -> Self: ... @overload def bfill( self, @@ -1495,7 +1600,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): inplace: Literal[False] = False, limit: int | None = ..., limit_area: Literal["inside", "outside"] | None = ..., - ) -> Series[S1]: ... + ) -> Self: ... @overload def interpolate( self, @@ -1519,7 +1624,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): limit_direction: Literal["forward", "backward", "both"] | None = ..., limit_area: Literal["inside", "outside"] | None = ..., **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... @final def asof( self, @@ -1555,7 +1660,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): axis: AxisIndex | None = 0, inplace: Literal[False] = False, **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... @final def asfreq( self, @@ -1564,14 +1669,14 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): how: Literal["start", "end"] | None = None, normalize: _bool = False, fill_value: Scalar | None = None, - ) -> Series[S1]: ... + ) -> Self: ... @final def at_time( self, time: _str | time, asof: _bool = False, axis: AxisIndex | None = 0, - ) -> Series[S1]: ... + ) -> Self: ... @final def between_time( self, @@ -1579,7 +1684,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): end_time: _str | time, inclusive: IntervalClosedType = "both", axis: AxisIndex | None = 0, - ) -> Series[S1]: ... + ) -> Self: ... @final def rank( self, @@ -1653,7 +1758,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): inplace: Literal[False] = False, axis: AxisIndex | None = 0, level: Level | None = ..., - ) -> Series[S1]: ... + ) -> Self: ... def case_when( self, caselist: list[ @@ -1672,7 +1777,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): after: date | _str | int | None = ..., axis: AxisIndex | None = 0, copy: _bool = ..., - ) -> Series[S1]: ... + ) -> Self: ... @final def tz_convert( self, @@ -1680,7 +1785,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): axis: AxisIndex = 0, level: Level | None = None, copy: _bool = True, - ) -> Series[S1]: ... + ) -> Self: ... @final def tz_localize( self, @@ -1690,16 +1795,16 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): copy: _bool = True, ambiguous: TimeAmbiguous = "raise", nonexistent: _str = "raise", - ) -> Series[S1]: ... + ) -> Self: ... @final - def abs(self) -> Series[S1]: ... + def abs(self) -> Self: ... @final def describe( self, percentiles: list[float] | None = ..., include: Literal["all"] | list[S1] | None = ..., exclude: S1 | list[S1] | None = ..., - ) -> Series[S1]: ... + ) -> Self: ... @final def pct_change( self, @@ -2953,9 +3058,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: num | ListLike | Series[S1]) -> Self: ... 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: num | ListLike | Series[S1]) -> Self: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] # pyrefly: ignore # bad-override @@ -2972,9 +3077,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: num | ListLike | Series[S1]) -> Self: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + def __rmod__(self, other: num | ListLike | Series[S1]) -> Self: ... + def __rpow__(self, other: num | ListLike | Series[S1]) -> Self: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] # pyrefly: ignore # bad-override @@ -4165,14 +4270,14 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): skipna: _bool = True, *args: Any, **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... def cummin( self, axis: AxisIndex | None = 0, skipna: _bool = True, *args: Any, **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... @overload def cumprod( self: Series[Never], @@ -4203,14 +4308,14 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): skipna: _bool = True, *args: Any, **kwargs: Any, - ) -> Series[S1]: ... + ) -> Self: ... def divmod( self, other: num | ListLike | Series[S1], level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex = ..., - ) -> Series[S1]: ... + ) -> Self: ... def eq( self, other: Scalar | Series[S1], @@ -4368,7 +4473,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex | None = 0, - ) -> Series[S1]: ... + ) -> Self: ... def ne( self, other: Scalar | Series[S1], @@ -4384,7 +4489,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex | None = 0, - ) -> Series[S1]: ... + ) -> Self: ... def prod( self, axis: AxisIndex | None = 0, @@ -4407,14 +4512,14 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex = ..., - ) -> Series[S1]: ... + ) -> Self: ... def rmod( self, other: Series[S1] | Scalar, level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex = ..., - ) -> Series[S1]: ... + ) -> Self: ... @overload def rolling( self, @@ -4447,7 +4552,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame): level: Level | None = ..., fill_value: float | None = None, axis: AxisIndex = ..., - ) -> Series[S1]: ... + ) -> Self: ... def sem( self, axis: AxisIndex | None = 0, diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 6c25d5e13..88a5f043a 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -12,10 +12,10 @@ 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 from pandas.core.arrays.datetimes import DatetimeArray from pandas.core.arrays.interval import IntervalArray +from pandas.core.arrays.numpy_ import NumpyExtensionArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.indexes.base import Index from pandas.core.indexes.category import CategoricalIndex @@ -1476,7 +1476,11 @@ def test_array_property() -> None: TimedeltaArray, pd.Timedelta, ) - check(assert_type(Index([1]).array, ExtensionArray), ExtensionArray, np.integer) + check( + assert_type(Index([1]).array, NumpyExtensionArray), + NumpyExtensionArray, + np.integer, + ) def test_to_series() -> None: diff --git a/tests/series/test_properties.py b/tests/series/test_properties.py index 21ce75931..2f61a2615 100644 --- a/tests/series/test_properties.py +++ b/tests/series/test_properties.py @@ -2,9 +2,9 @@ import numpy as np from pandas.core.arrays import DatetimeArray -from pandas.core.arrays.base import ExtensionArray from pandas.core.arrays.categorical import Categorical from pandas.core.arrays.interval import IntervalArray +from pandas.core.arrays.numpy_ import NumpyExtensionArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.frame import DataFrame from pandas.core.indexes.accessors import ( @@ -77,7 +77,9 @@ def test_property_array() -> None: TimedeltaArray, Timedelta, ) - check(assert_type(Series([1]).array, ExtensionArray), ExtensionArray, np.integer) - # python/mypy#19952: mypy believes ExtensionArray and its subclasses have a - # conflict and gives Any for s.array - check(assert_type(Series([1, "s"]).array, ExtensionArray), ExtensionArray) # type: ignore[assert-type] + check( + assert_type(Series([1]).array, NumpyExtensionArray), + NumpyExtensionArray, + np.integer, + ) + check(assert_type(Series([1, "s"]).array, NumpyExtensionArray), NumpyExtensionArray) diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 73bc51bb5..1b940312c 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -22,16 +22,12 @@ TypeAlias, TypedDict, TypeVar, - cast, ) import numpy as np import pandas as pd -from pandas.api.extensions import ( - ExtensionArray, - ExtensionDtype, -) from pandas.api.typing import NAType +from pandas.core import arrays from pandas.core.arrays.datetimes import DatetimeArray from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.window import ExponentialMovingWindow @@ -50,7 +46,7 @@ Scalar, ) -from pandas.core.dtypes.dtypes import CategoricalDtype # noqa F401 +from pandas.core.dtypes.base import ExtensionDtype from tests import ( PD_LTE_23, @@ -1412,29 +1408,29 @@ def test_types_rename_axis() -> None: def test_types_values() -> None: check( - assert_type(pd.Series([1, 2, 3]).values, ExtensionArray | np.ndarray), + assert_type(pd.Series([1, 2, 3]).values, arrays.ExtensionArray | np.ndarray), np.ndarray, ) - valresult_type: type[np.ndarray | ExtensionArray] + valresult_type: type[np.ndarray | arrays.ExtensionArray] if PD_LTE_23: valresult_type = np.ndarray else: - valresult_type = ExtensionArray + valresult_type = arrays.ExtensionArray check( - assert_type(pd.Series(list("aabc")).values, np.ndarray | ExtensionArray), + assert_type(pd.Series(list("aabc")).values, np.ndarray | arrays.ExtensionArray), valresult_type, ) check( assert_type( pd.Series(list("aabc")).astype("category").values, - np.ndarray | ExtensionArray, + np.ndarray | arrays.ExtensionArray, ), pd.Categorical, ) check( assert_type( pd.Series(pd.date_range("20130101", periods=3, tz="US/Eastern")).values, - np.ndarray | ExtensionArray, + np.ndarray | arrays.ExtensionArray, ), np.ndarray, ) @@ -1697,7 +1693,7 @@ def test_series_replace() -> None: def test_cat_accessor() -> None: # GH 43 - s: pd.Series[str] = pd.Series( + s: pd.Series[str, pd.Categorical] = pd.Series( pd.Categorical(["a", "b", "a"], categories=["a", "b"]) ) check(assert_type(s.cat.codes, "pd.Series[int]"), pd.Series, np.int8) @@ -1830,7 +1826,7 @@ def test_categorical_codes() -> None: # GH1383 sr = pd.Series([1], dtype="category") - check(assert_type(sr, "pd.Series[CategoricalDtype]"), pd.Series, np.integer) + check(assert_type(sr, "pd.Series[int, pd.Categorical]"), pd.Series, np.integer) def test_relops() -> None: @@ -2916,10 +2912,17 @@ def test_astype_categorical(cast_arg: CategoryDtypeArg, target_type: type) -> No s = pd.Series(["a", "b"]) check(s.astype(cast_arg), pd.Series, target_type) - if TYPE_CHECKING: - # pandas category - assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[pd.CategoricalDtype]") - assert_type(s.astype(cast_arg), "pd.Series[pd.CategoricalDtype]") + # pandas category + check( + assert_type(s.astype(pd.CategoricalDtype()), "pd.Series[str, pd.Categorical]"), + pd.Series, + str, + ) + check( + assert_type(s.astype(cast_arg), "pd.Series[str, pd.Categorical]"), + pd.Series, + str, + ) @pytest.mark.parametrize("cast_arg, target_type", ASTYPE_OBJECT_ARGS, ids=repr) @@ -3497,14 +3500,10 @@ def test_diff() -> None: index_to_check_for_type=-1, ) # nullable bool -> nullable bool - # casting due to pandas-dev/pandas-stubs#1395 check( assert_type( - cast( - "pd.Series[pd.BooleanDtype]", - pd.Series([True, True, False, False, True], dtype="boolean").diff(), - ), - "pd.Series[pd.BooleanDtype]", + pd.Series([True, True, False, False, True], dtype="boolean").diff(), + "pd.Series[bool, arrays.BooleanArray]", ), pd.Series, np.bool_, diff --git a/tests/test_natype.py b/tests/test_natype.py index 7a524bc37..a87cf2ed8 100644 --- a/tests/test_natype.py +++ b/tests/test_natype.py @@ -2,7 +2,7 @@ import pandas as pd from pandas.api.typing import NAType -from pandas.core.arrays.boolean import BooleanArray +from pandas.core import arrays from typing_extensions import assert_type from tests import check @@ -15,32 +15,32 @@ def test_arithmetic() -> None: idx_int: pd.Index[int] = pd.Index([1, 2, 3], dtype="Int64") # __add__ - check(assert_type(na + s_int, pd.Series), pd.Series) + check(assert_type(na + s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na + idx_int, pd.Index), pd.Index) check(assert_type(na + 1, NAType), NAType) # __radd__ - check(assert_type(s_int + na, pd.Series), pd.Series) + check(assert_type(s_int + na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int + na, pd.Index), pd.Index) check(assert_type(1 + na, NAType), NAType) # __sub__ - check(assert_type(na - s_int, pd.Series), pd.Series) + check(assert_type(na - s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na - idx_int, pd.Index), pd.Index) check(assert_type(na - 1, NAType), NAType) # __rsub__ - check(assert_type(s_int - na, pd.Series), pd.Series) + check(assert_type(s_int - na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int - na, pd.Index), pd.Index) check(assert_type(1 - na, NAType), NAType) # __mul__ - check(assert_type(na * s_int, pd.Series), pd.Series) + check(assert_type(na * s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na * idx_int, pd.Index), pd.Index) check(assert_type(na * 1, NAType), NAType) # __rmul__ - check(assert_type(s_int * na, pd.Series), pd.Series) + check(assert_type(s_int * na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int * na, pd.Index), pd.Index) check(assert_type(1 * na, NAType), NAType) @@ -51,32 +51,32 @@ def test_arithmetic() -> None: check(assert_type(1 @ na, NAType), NAType) # __truediv__ - check(assert_type(na / s_int, pd.Series), pd.Series) + check(assert_type(na / s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na / idx_int, pd.Index), pd.Index) check(assert_type(na / 1, NAType), NAType) # __rtruediv__ - check(assert_type(s_int / na, pd.Series), pd.Series) + check(assert_type(s_int / na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int / na, pd.Index), pd.Index) check(assert_type(1 / na, NAType), NAType) # __floordiv__ - check(assert_type(na // s_int, pd.Series), pd.Series) + check(assert_type(na // s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na // idx_int, pd.Index), pd.Index) check(assert_type(na // 1, NAType), NAType) # __rfloordiv__ - check(assert_type(s_int // na, pd.Series), pd.Series) + check(assert_type(s_int // na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int // na, pd.Index), pd.Index) check(assert_type(1 // na, NAType), NAType) # __mod__ - check(assert_type(na % s_int, pd.Series), pd.Series) + check(assert_type(na % s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na % idx_int, pd.Index), pd.Index) check(assert_type(na % 1, NAType), NAType) # __rmod__ - check(assert_type(s_int % na, pd.Series), pd.Series) + check(assert_type(s_int % na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int % na, "pd.Index[int]"), pd.Index) check(assert_type(1 % na, NAType), NAType) @@ -121,42 +121,42 @@ def test_arithmetic() -> None: check(assert_type(divmod(1, na), tuple[NAType, NAType]), tuple) # __eq__ - check(assert_type(na == s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na == idx_int, BooleanArray), BooleanArray) + check(assert_type(na == s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na == idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na == 1, NAType), NAType) # __ne__ - check(assert_type(na != s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na != idx_int, BooleanArray), BooleanArray) + check(assert_type(na != s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na != idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na != 1, NAType), NAType) # __le__ - check(assert_type(na <= s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na <= idx_int, BooleanArray), BooleanArray) + check(assert_type(na <= s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na <= idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na <= 1, NAType), NAType) # __lt__ - check(assert_type(na < s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na < idx_int, BooleanArray), BooleanArray) + check(assert_type(na < s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na < idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na < 1, NAType), NAType) # __gt__ - check(assert_type(na > s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na > idx_int, BooleanArray), BooleanArray) + check(assert_type(na > s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na > idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na > 1, NAType), NAType) # __ge__ - check(assert_type(na >= s_int, "pd.Series[bool]"), pd.Series) - check(assert_type(na >= idx_int, BooleanArray), BooleanArray) + check(assert_type(na >= s_int, "pd.Series[bool, arrays.BooleanArray]"), pd.Series) + check(assert_type(na >= idx_int, arrays.BooleanArray), arrays.BooleanArray) check(assert_type(na >= 1, NAType), NAType) # __pow__ - check(assert_type(na**s_int, pd.Series), pd.Series) + check(assert_type(na**s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na**idx_int, pd.Index), pd.Index) check(assert_type(na**2, NAType), NAType) # __rpow__ - check(assert_type(s_int**na, pd.Series), pd.Series) + check(assert_type(s_int**na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int**na, "pd.Index[int]"), pd.Index) check(assert_type(2**na, NAType), NAType) @@ -178,9 +178,9 @@ def test_arithmetic() -> None: check(assert_type(True | na, Literal[True]), bool) # __xor__ - check(assert_type(na ^ s_int, pd.Series), pd.Series) + check(assert_type(na ^ s_int, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(na ^ idx_int, pd.Index), pd.Index) # rxor - check(assert_type(s_int ^ na, pd.Series), pd.Series) + check(assert_type(s_int ^ na, "pd.Series[int, arrays.IntegerArray]"), pd.Series) check(assert_type(idx_int ^ na, pd.Index), pd.Index) diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 6925582ac..dc5c6c4ee 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -1025,9 +1025,21 @@ def test_cut() -> None: 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") i = pd.cut(pd.Series([1, 2, 3, 4, 5, 6, 7, 8]), 4, labels=["1", "2", "3", "4"]) - check(assert_type(g, pd.Series), pd.Series) - check(assert_type(h, pd.Series), pd.Series) - check(assert_type(i, pd.Series), pd.Series) + check( + assert_type(g, "pd.Series[pd.Interval[int], pd.Categorical]"), + pd.Series, + pd.Interval, + ) + check( + assert_type(h, "pd.Series[pd.Interval[int], pd.Categorical]"), + pd.Series, + pd.Interval, + ) + check( + assert_type(i, "pd.Series[pd.Interval[int], pd.Categorical]"), + pd.Series, + pd.Interval, + ) j0, j1 = pd.cut( pd.Series([1, 2, 3, 4, 5, 6, 7, 8]), @@ -1071,14 +1083,14 @@ def test_cut() -> None: check( assert_type( pd.cut(s1, bins=[np.datetime64("2020-01-03"), np.datetime64("2020-09-01")]), - "pd.Series[pd.CategoricalDtype]", + "pd.Series[pd.Interval[pd.Timestamp], pd.Categorical]", ), pd.Series, ) check( assert_type( pd.cut(s1, bins=10), - "pd.Series[pd.CategoricalDtype]", + "pd.Series[pd.Interval[pd.Timestamp], pd.Categorical]", ), pd.Series, pd.Interval,