From 4929ecb497c1d66a1ea8489d01bbaaf29dec0f4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Wed, 2 Aug 2023 09:10:58 -0400 Subject: [PATCH 01/23] interpolate method --- .pre-commit-config.yaml | 2 +- pandas-stubs/_typing.pyi | 20 ++++++++++++++++++++ pandas-stubs/core/frame.pyi | 7 ++++--- pandas-stubs/core/resample.pyi | 27 +++------------------------ pandas-stubs/core/series.pyi | 23 ++--------------------- tests/test_scalars.py | 2 +- 6 files changed, 31 insertions(+), 50 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cc0e9aaa0..1786d3276 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,7 @@ repos: hooks: - id: isort - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.278 + rev: v0.0.282 hooks: - id: ruff args: [ diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 27548758f..4306aede0 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -664,6 +664,26 @@ StataDateFormat: TypeAlias = Literal[ ] FillnaOptions: TypeAlias = Literal["backfill", "bfill", "ffill", "pad"] +InterpolateOptions: TypeAlias = Literal[ + "linear", + "time", + "index", + "pad", + "nearest", + "zero", + "slinear", + "quadratic", + "cubic", + "barycentric", + "polynomial", + "krogh", + "piecewise_polynomial", + "spline", + "pchip", + "akima", + "cubicspline", + "from_derivatives", +] ReplaceMethod: TypeAlias = Literal["pad", "ffill", "bfill"] SortKind: TypeAlias = Literal["quicksort", "mergesort", "heapsort", "stable"] NaPosition: TypeAlias = Literal["first", "last"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index f9c195b18..a8f46facb 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -86,6 +86,7 @@ from pandas._typing import ( IndexingInt, IndexLabel, IndexType, + InterpolateOptions, IntervalClosedType, IntervalT, JoinHow, @@ -1704,7 +1705,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: _str = ..., + method: InterpolateOptions = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -1717,7 +1718,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: _str = ..., + method: InterpolateOptions = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -1730,7 +1731,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: _str = ..., + method: InterpolateOptions = ..., *, axis: Axis = ..., limit: int | None = ..., diff --git a/pandas-stubs/core/resample.pyi b/pandas-stubs/core/resample.pyi index 23de00bf3..a21b7dfcf 100644 --- a/pandas-stubs/core/resample.pyi +++ b/pandas-stubs/core/resample.pyi @@ -22,6 +22,7 @@ from typing_extensions import TypeAlias from pandas._typing import ( Axis, + InterpolateOptions, NDFrameT, Scalar, npt, @@ -50,28 +51,6 @@ _SeriesGroupByFuncArgs: TypeAlias = ( _SeriesGroupByFuncTypes | Mapping[Hashable, _SeriesGroupByFunc | str] ) -_Interpolation: TypeAlias = Literal[ - "linear", - "time", - "index", - "pad", - "nearest", - "zero", - "slinear", - "quadratic", - "cubic", - "spline", - "barycentric", - "polynomial", - "krogh", - "piecewise_polynomial", - "spline", - "pchip", - "akima", - "cubicspline", - "from_derivatives", -] - class Resampler(BaseGroupBy, Generic[NDFrameT]): def __getattr__(self, attr: str) -> SeriesGroupBy: ... def __iter__(self) -> Generator[tuple[Hashable, NDFrameT], None, None]: ... @@ -147,7 +126,7 @@ class Resampler(BaseGroupBy, Generic[NDFrameT]): @overload def interpolate( self, - method: _Interpolation = ..., + method: InterpolateOptions = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -160,7 +139,7 @@ class Resampler(BaseGroupBy, Generic[NDFrameT]): @overload def interpolate( self, - method: _Interpolation = ..., + method: InterpolateOptions = ..., *, axis: Axis = ..., limit: int | None = ..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c05e74237..813491e0f 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -118,6 +118,7 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntDtypeArg, + InterpolateOptions, IntervalClosedType, IntervalT, JoinHow, @@ -1273,27 +1274,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): ) -> Series[S1] | None: ... def interpolate( self, - method: _str - | Literal[ - "linear", - "time", - "index", - "values", - "pad", - "nearest", - "slinear", - "quadratic", - "cubic", - "spline", - "barycentric", - "polynomial", - "krogh", - "pecewise_polynomial", - "spline", - "pchip", - "akima", - "from_derivatives", - ] = ..., + method: InterpolateOptions = ..., *, axis: AxisIndex | None = ..., limit: int | None = ..., diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 958a03b61..8c946839d 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -1745,7 +1745,7 @@ def test_period_add_subtract() -> None: as_dt_td = dt.timedelta(days=1) as_np_td = np.timedelta64(1, "D") as_np_i64 = np.int64(1) - as_int = int(1) + as_int: int = 1 as_period_index = pd.period_range("2012-1-1", periods=10, freq="D") check(assert_type(as_period_index, pd.PeriodIndex), pd.PeriodIndex) as_period = pd.Period("2012-1-1", freq="D") From 2775a58547b115a58f527d236f69148e3a2cae32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Wed, 2 Aug 2023 11:00:14 -0400 Subject: [PATCH 02/23] WIP --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 17 +-- pandas-stubs/core/base.pyi | 21 +-- pandas-stubs/core/indexes/accessors.pyi | 11 +- pandas-stubs/core/indexes/base.pyi | 99 +++++++------ pandas-stubs/core/indexes/datetimelike.pyi | 9 +- pandas-stubs/core/indexes/datetimes.pyi | 14 +- pandas-stubs/core/indexes/extension.pyi | 4 +- pandas-stubs/core/indexes/interval.pyi | 11 +- pandas-stubs/core/indexes/period.pyi | 17 +-- pandas-stubs/core/indexes/range.pyi | 9 +- pandas-stubs/core/indexes/timedeltas.pyi | 15 +- pandas-stubs/core/reshape/tile.pyi | 58 +++----- pandas-stubs/core/series.pyi | 2 +- tests/test_frame.py | 2 +- tests/test_indexes.py | 153 +++++++++++---------- tests/test_pandas.py | 2 +- tests/test_timefuncs.py | 67 +++++---- 17 files changed, 230 insertions(+), 281 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 1ded708c9..7f5652073 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -11,14 +11,11 @@ import numpy as np import pandas as pd from pandas import ( DatetimeIndex, + Index, PeriodIndex, Series, TimedeltaIndex, ) -from pandas.core.indexes.base import ( - _FloatIndexType, - _IntIndexType, -) from pandas.core.series import ( TimedeltaSeries, TimestampSeries, @@ -245,7 +242,7 @@ class Timedelta(timedelta): @overload def __mul__(self, other: Series[float]) -> TimedeltaSeries: ... @overload - def __mul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ... + def __mul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... @overload def __rmul__(self, other: float) -> Timedelta: ... @overload @@ -255,7 +252,7 @@ class Timedelta(timedelta): @overload def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ... @overload - def __rmul__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ... + def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... # Override due to more types supported than dt.timedelta # error: Signature of "__floordiv__" incompatible with supertype "timedelta" @overload # type: ignore[override] @@ -271,9 +268,7 @@ class Timedelta(timedelta): self, other: npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.int_]: ... @overload - def __floordiv__( - self, other: _IntIndexType | _FloatIndexType - ) -> TimedeltaIndex: ... + def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... @overload def __floordiv__(self, other: Series[int]) -> TimedeltaSeries: ... @overload @@ -306,7 +301,7 @@ class Timedelta(timedelta): @overload def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ... @overload - def __truediv__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ... + def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ... # Override due to more types supported than dt.timedelta @overload @@ -338,7 +333,7 @@ class Timedelta(timedelta): @overload def __mod__(self, other: Series[int] | Series[float]) -> TimedeltaSeries: ... @overload - def __mod__(self, other: _IntIndexType | _FloatIndexType) -> TimedeltaIndex: ... + def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... @overload def __mod__( self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index eeec4b511..865d59e7a 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -1,5 +1,6 @@ from typing import ( Generic, + Iterator, Literal, ) @@ -9,8 +10,10 @@ from pandas.core.accessor import DirNamesMixin from pandas.core.arraylike import OpsMixin from pandas.core.arrays import ExtensionArray from pandas.core.arrays.categorical import Categorical +from typing_extensions import Self from pandas._typing import ( + S1, AxisIndex, NaPosition, NDFrameT, @@ -28,16 +31,16 @@ class SelectionMixin(Generic[NDFrameT]): def ndim(self) -> int: ... def __getitem__(self, key): ... -class IndexOpsMixin(OpsMixin): +class IndexOpsMixin(OpsMixin, Generic[S1]): __array_priority__: int = ... - def transpose(self, *args, **kwargs) -> IndexOpsMixin: ... + def transpose(self, *args, **kwargs) -> Self: ... @property - def T(self) -> IndexOpsMixin: ... + def T(self) -> Self: ... @property def shape(self) -> tuple: ... @property def ndim(self) -> int: ... - def item(self): ... + def item(self) -> S1: ... @property def nbytes(self) -> int: ... @property @@ -61,9 +64,9 @@ class IndexOpsMixin(OpsMixin): def argmin( self, axis: AxisIndex | None = ..., skipna: bool = ..., *args, **kwargs ) -> np.int64: ... - def tolist(self) -> list: ... - def to_list(self) -> list: ... - def __iter__(self): ... + def tolist(self) -> list[S1]: ... + def to_list(self) -> list[S1]: ... + def __iter__(self) -> Iterator[S1]: ... @property def hasnans(self) -> bool: ... def value_counts( @@ -87,6 +90,4 @@ class IndexOpsMixin(OpsMixin): def searchsorted( self, value, side: Literal["left", "right"] = ..., sorter=... ) -> int | list[int]: ... - def drop_duplicates( - self, *, keep: NaPosition | Literal[False] = ... - ) -> IndexOpsMixin: ... + def drop_duplicates(self, *, keep: NaPosition | Literal[False] = ...) -> Self: ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index 699a74a30..f181f7036 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -25,7 +25,6 @@ from pandas.core.base import ( PandasObject, ) from pandas.core.frame import DataFrame -from pandas.core.indexes.base import _IntIndexType from pandas.core.series import ( PeriodSeries, Series, @@ -43,7 +42,7 @@ from pandas._typing import ( class Properties(PandasDelegate, PandasObject, NoNewAttributesMixin): def __init__(self, data: Series, orig) -> None: ... -_DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", Series[int], _IntIndexType) +_DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", Series[int], Index[int]) class _DayLikeFieldOps(Generic[_DTFieldOpsReturnType]): @property @@ -302,7 +301,7 @@ class TimedeltaProperties( ): ... _PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex) -_PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], _IntIndexType) +_PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int]) _PeriodStrReturnTypes = TypeVar("_PeriodStrReturnTypes", Series[str], Index) _PeriodDTAReturnTypes = TypeVar("_PeriodDTAReturnTypes", DatetimeArray, DatetimeIndex) _PeriodPAReturnTypes = TypeVar("_PeriodPAReturnTypes", PeriodArray, PeriodIndex) @@ -335,8 +334,8 @@ class _PeriodProperties( ) -> _PeriodPAReturnTypes: ... class PeriodIndexFieldOps( - _DayLikeFieldOps[_IntIndexType], - _PeriodProperties[DatetimeIndex, _IntIndexType, Index, DatetimeIndex, PeriodIndex], + _DayLikeFieldOps[Index[int]], + _PeriodProperties[DatetimeIndex, Index[int], Index, DatetimeIndex, PeriodIndex], ): ... class PeriodProperties( Properties, @@ -382,7 +381,7 @@ class TimestampProperties( class DatetimeIndexProperties( Properties, _DatetimeNoTZProperties[ - _IntIndexType, + Index[int], np_ndarray_bool, DatetimeIndex, np.ndarray, diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index b9ad3a1f7..106009cd9 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -5,9 +5,10 @@ from collections.abc import ( Iterator, Sequence, ) +from datetime import timedelta from typing import ( + Any, ClassVar, - Generic, Literal, overload, ) @@ -17,6 +18,7 @@ from pandas import ( DataFrame, MultiIndex, Series, + TimedeltaIndex, ) from pandas.core.arrays import ExtensionArray from pandas.core.base import ( @@ -37,7 +39,6 @@ from pandas._typing import ( DtypeObj, FillnaOptions, HashableT, - IndexT, Label, Level, NaPosition, @@ -52,22 +53,7 @@ class InvalidIndexError(Exception): ... _str = str -class _IndexGetitemMixin(Generic[S1]): - @overload - def __getitem__( - self, - idx: slice - | np_ndarray_anyint - | Sequence[int] - | Index - | Series[bool] - | Sequence[bool] - | np_ndarray_bool, - ) -> Self: ... - @overload - def __getitem__(self, idx: int) -> S1: ... - -class Index(IndexOpsMixin, PandasObject): +class Index(IndexOpsMixin[S1], PandasObject): __hash__: ClassVar[None] # type: ignore[assignment] @overload def __new__( # type: ignore[misc] @@ -78,7 +64,7 @@ class Index(IndexOpsMixin, PandasObject): name=..., tupleize_cols: bool = ..., **kwargs, - ) -> _IntIndexType: ... + ) -> Index[int]: ... @overload def __new__( # type: ignore[misc] cls, @@ -91,7 +77,7 @@ class Index(IndexOpsMixin, PandasObject): name=..., tupleize_cols: bool = ..., **kwargs, - ) -> _FloatIndexType: ... + ) -> Index[float]: ... @overload def __new__( # type: ignore[misc] cls, @@ -101,19 +87,29 @@ class Index(IndexOpsMixin, PandasObject): name=..., tupleize_cols: bool = ..., **kwargs, - ) -> _ComplexIndexType: ... + ) -> Index[complex]: ... @overload def __new__( cls, - data: Iterable = ..., + data: Iterable[S1] = ..., dtype=..., copy: bool = ..., name=..., tupleize_cols: bool = ..., **kwargs, ) -> Self: ... + @overload + def __new__( + cls, + data: Iterable = ..., + dtype=..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> Index: ... @property - def str(self) -> StringMethods[Index, MultiIndex]: ... + def str(self) -> StringMethods[Self, MultiIndex]: ... @property def asi8(self) -> np_ndarray_int64: ... def is_(self, other) -> bool: ... @@ -129,7 +125,7 @@ class Index(IndexOpsMixin, PandasObject): self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs ): ... def repeat(self, repeats, axis=...): ... - def copy(self, name=..., deep: bool = ...) -> Index: ... + def copy(self, name=..., deep: bool = ...) -> Self: ... def __copy__(self, **kwargs): ... def __deepcopy__(self, memo=...): ... def format( @@ -181,11 +177,9 @@ class Index(IndexOpsMixin, PandasObject): def notna(self): ... notnull = ... def fillna(self, value=..., downcast=...): ... - def dropna(self, how: Literal["any", "all"] = ...) -> Index: ... - def unique(self, level=...) -> Index: ... - def drop_duplicates( - self, *, keep: NaPosition | Literal[False] = ... - ) -> IndexOpsMixin: ... + def dropna(self, how: Literal["any", "all"] = ...) -> Self: ... + def unique(self, level=...) -> Self: ... + def drop_duplicates(self, *, keep: NaPosition | Literal[False] = ...) -> Self: ... def duplicated( self, keep: Literal["first", "last", False] = ... ) -> np_ndarray_bool: ... @@ -195,15 +189,15 @@ class Index(IndexOpsMixin, PandasObject): def __ror__(self, other: Never) -> Never: ... def __xor__(self, other: Never) -> Never: ... def __rxor__(self, other: Never) -> Never: ... - def __neg__(self: IndexT) -> IndexT: ... + def __neg__(self) -> Self: ... def __nonzero__(self) -> None: ... __bool__ = ... def union(self, other: list[HashableT] | Index, sort=...) -> Index: ... - def intersection(self, other: list[T1] | Index, sort: bool = ...) -> Index: ... - def difference(self, other: list | Index, sort: bool | None = None) -> Index: ... + def intersection(self, other: list[T1] | Self, sort: bool = ...) -> Self: ... + def difference(self, other: list | Index, sort: bool | None = None) -> Self: ... def symmetric_difference( self, other: list[T1] | Index, result_name=..., sort=... - ) -> Index: ... + ) -> Self: ... def get_loc( self, key: Label, @@ -232,7 +226,7 @@ class Index(IndexOpsMixin, PandasObject): def __setitem__(self, key, value) -> None: ... @overload def __getitem__( - self: IndexT, + self, idx: slice | np_ndarray_anyint | Sequence[int] @@ -240,9 +234,9 @@ class Index(IndexOpsMixin, PandasObject): | Series[bool] | Sequence[bool] | np_ndarray_bool, - ) -> IndexT: ... + ) -> Self: ... @overload - def __getitem__(self, idx: int | tuple[np_ndarray_anyint, ...]) -> Scalar: ... + def __getitem__(self, idx: int | tuple[np_ndarray_anyint, ...]) -> S1: ... def append(self, other): ... def putmask(self, mask, value): ... def equals(self, other) -> bool: ... @@ -265,7 +259,7 @@ class Index(IndexOpsMixin, PandasObject): def slice_locs(self, start=..., end=..., step=...): ... def delete(self, loc): ... def insert(self, loc, item): ... - def drop(self, labels, errors: _str = ...) -> Index: ... + def drop(self, labels, errors: _str = ...) -> Self: ... @property def shape(self) -> tuple[int, ...]: ... # Extra methods from old stubs @@ -276,25 +270,24 @@ class Index(IndexOpsMixin, PandasObject): def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] def __lt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] def __gt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] + # overwrite inherit methods from OpsMixin + @overload + def __mul__( # type: ignore[misc] + self: Index[int] | Index[float], other: timedelta + ) -> TimedeltaIndex: ... + @overload + def __mul__(self, other: Any) -> Self: ... + @overload + def __floordiv__(self: Index[int] | Index[float], other: timedelta) -> Never: ... + @overload + def __floordiv__(self, other: Any) -> Self: ... + @overload + def __truediv__(self: Index[int] | Index[float], other: timedelta) -> Never: ... + @overload + def __truediv__(self, other: Any) -> Self: ... def ensure_index_from_sequences( sequences: Sequence[Sequence[Dtype]], names: list[str] = ... ) -> Index: ... def ensure_index(index_like: Sequence | Index, copy: bool = ...) -> Index: ... def maybe_extract_name(name, obj, cls) -> Label: ... - -class _NumericIndexType(Index): - def __add__(self, other: _NumericIndexType | complex) -> Self: ... - def __radd__(self, other: _NumericIndexType | complex) -> Self: ... - def __sub__(self, other: _NumericIndexType | complex) -> Self: ... - def __rsub__(self, other: _NumericIndexType | complex) -> Self: ... - def __mul__(self, other: _NumericIndexType | complex) -> Self: ... - def __rmul__(self, other: _NumericIndexType | complex) -> Self: ... - def __truediv__(self, other: _NumericIndexType | complex) -> Self: ... - def __rtruediv__(self, other: _NumericIndexType | complex) -> Self: ... - def __floordiv__(self, other: _NumericIndexType | complex) -> Self: ... - def __rfloordiv__(self, other: _NumericIndexType | complex) -> Self: ... - -class _FloatIndexType(_NumericIndexType): ... -class _IntIndexType(_NumericIndexType): ... -class _ComplexIndexType(_NumericIndexType): ... diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index e0425bd95..6e808f900 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -2,8 +2,9 @@ from pandas.core.indexes.extension import ExtensionIndex from pandas.core.indexes.timedeltas import TimedeltaIndex from pandas._libs.tslibs import BaseOffset +from pandas._typing import S1 -class DatetimeIndexOpsMixin(ExtensionIndex): +class DatetimeIndexOpsMixin(ExtensionIndex[S1]): @property def freq(self) -> BaseOffset | None: ... @property @@ -14,6 +15,8 @@ class DatetimeIndexOpsMixin(ExtensionIndex): def argmin(self, axis=..., skipna: bool = ..., *args, **kwargs): ... def max(self, axis=..., skipna: bool = ..., *args, **kwargs): ... def argmax(self, axis=..., skipna: bool = ..., *args, **kwargs): ... - def __rsub__(self, other: DatetimeIndexOpsMixin) -> TimedeltaIndex: ... + def __rsub__( # type: ignore[override] + self, other: DatetimeIndexOpsMixin + ) -> TimedeltaIndex: ... -class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin): ... +class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): ... diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index baa06efdf..9ece274a9 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -14,15 +14,12 @@ from typing import ( import numpy as np from pandas import ( DataFrame, + Index, Timedelta, TimedeltaIndex, Timestamp, ) from pandas.core.indexes.accessors import DatetimeIndexProperties -from pandas.core.indexes.base import ( - _FloatIndexType, - _IndexGetitemMixin, -) from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin from pandas.core.series import ( TimedeltaSeries, @@ -40,12 +37,7 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype from pandas.tseries.offsets import BaseOffset -# type ignore needed because of __getitem__() -class DatetimeIndex( # type: ignore[misc] - _IndexGetitemMixin[Timestamp], - DatetimeTimedeltaMixin, - DatetimeIndexProperties, -): +class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties): def __init__( self, data: ArrayLike | AnyArrayLike | list | tuple, @@ -91,7 +83,7 @@ class DatetimeIndex( # type: ignore[misc] self, start_time, end_time, include_start: bool = ..., include_end: bool = ... ): ... def to_perioddelta(self, freq) -> TimedeltaIndex: ... - def to_julian_date(self) -> _FloatIndexType: ... + def to_julian_date(self) -> Index[float]: ... def isocalendar(self) -> DataFrame: ... @property def tzinfo(self) -> tzinfo | None: ... diff --git a/pandas-stubs/core/indexes/extension.pyi b/pandas-stubs/core/indexes/extension.pyi index a5595f622..293784afa 100644 --- a/pandas-stubs/core/indexes/extension.pyi +++ b/pandas-stubs/core/indexes/extension.pyi @@ -1,3 +1,5 @@ from pandas.core.indexes.base import Index -class ExtensionIndex(Index): ... +from pandas._typing import S1 + +class ExtensionIndex(Index[S1]): ... diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index aeef461fe..6b45ebca0 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -4,7 +4,6 @@ from collections.abc import ( ) import datetime as dt from typing import ( - Generic, Literal, overload, ) @@ -12,10 +11,6 @@ from typing import ( import numpy as np import pandas as pd from pandas import Index -from pandas.core.indexes.base import ( - _FloatIndexType, - _IntIndexType, -) from pandas.core.indexes.extension import ExtensionIndex from pandas.core.series import ( Series, @@ -50,10 +45,10 @@ _EdgesInt: TypeAlias = ( | npt.NDArray[np.int32] | npt.NDArray[np.intp] | pd.Series[int] - | _IntIndexType + | Index[int] ) _EdgesFloat: TypeAlias = ( - Sequence[float] | npt.NDArray[np.float64] | pd.Series[float] | _FloatIndexType + Sequence[float] | npt.NDArray[np.float64] | pd.Series[float] | Index[float] ) _EdgesTimestamp: TypeAlias = ( Sequence[DatetimeLike] @@ -70,7 +65,7 @@ _EdgesTimedelta: TypeAlias = ( _TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime _TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta -class IntervalIndex(ExtensionIndex, IntervalMixin, Generic[IntervalT]): +class IntervalIndex(ExtensionIndex[IntervalT], IntervalMixin): closed: IntervalClosedType def __new__( diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 49b23b102..9580dda72 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/pandas-stubs/core/indexes/period.pyi @@ -6,11 +6,11 @@ import numpy as np import pandas as pd from pandas import Index from pandas.core.indexes.accessors import PeriodIndexFieldOps -from pandas.core.indexes.base import _IndexGetitemMixin from pandas.core.indexes.datetimelike import ( DatetimeIndexOpsMixin as DatetimeIndexOpsMixin, ) from pandas.core.indexes.timedeltas import TimedeltaIndex +from typing_extensions import Self from pandas._libs.tslibs import ( BaseOffset, @@ -19,12 +19,7 @@ from pandas._libs.tslibs import ( ) from pandas._libs.tslibs.period import _PeriodAddSub -# type ignore needed because of __getitem__() -class PeriodIndex( # type: ignore[misc] - _IndexGetitemMixin[Period], - DatetimeIndexOpsMixin, - PeriodIndexFieldOps, -): +class PeriodIndex(DatetimeIndexOpsMixin[pd.Period], PeriodIndexFieldOps): def __new__( cls, data=..., @@ -42,17 +37,17 @@ class PeriodIndex( # type: ignore[misc] @overload def __sub__(self, other: Period) -> Index: ... @overload - def __sub__(self, other: PeriodIndex) -> Index: ... + def __sub__(self, other: Self) -> Index: ... @overload - def __sub__(self, other: _PeriodAddSub) -> PeriodIndex: ... + def __sub__(self, other: _PeriodAddSub) -> Self: ... @overload def __sub__(self, other: NaTType) -> NaTType: ... @overload - def __sub__(self, other: TimedeltaIndex | pd.Timedelta) -> PeriodIndex: ... + def __sub__(self, other: TimedeltaIndex | pd.Timedelta) -> Self: ... @overload # type: ignore[override] def __rsub__(self, other: Period) -> Index: ... @overload - def __rsub__(self, other: PeriodIndex) -> Index: ... + def __rsub__(self, other: Self) -> Index: ... @overload def __rsub__(self, other: NaTType) -> NaTType: ... def __array__(self, dtype=...) -> np.ndarray: ... diff --git a/pandas-stubs/core/indexes/range.pyi b/pandas-stubs/core/indexes/range.pyi index 3166fc7d2..209b07d8c 100644 --- a/pandas-stubs/core/indexes/range.pyi +++ b/pandas-stubs/core/indexes/range.pyi @@ -3,10 +3,7 @@ from typing import overload import numpy as np from pandas import Series -from pandas.core.indexes.base import ( - Index, - _IntIndexType, -) +from pandas.core.indexes.base import Index from pandas._typing import ( HashableT, @@ -15,7 +12,7 @@ from pandas._typing import ( npt, ) -class RangeIndex(_IntIndexType): +class RangeIndex(Index[int]): def __new__( cls, start: int | RangeIndex = ..., @@ -83,7 +80,7 @@ class RangeIndex(_IntIndexType): def any(self) -> bool: ... def union( self, other: list[HashableT] | Index, sort=... - ) -> Index | _IntIndexType | RangeIndex: ... + ) -> Index | Index[int] | RangeIndex: ... @overload # type: ignore[override] def __getitem__( self, diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index e4b7fd02b..076bf767c 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -14,11 +14,11 @@ from pandas import ( Period, ) from pandas.core.indexes.accessors import TimedeltaIndexProperties -from pandas.core.indexes.base import _IndexGetitemMixin from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin from pandas.core.indexes.datetimes import DatetimeIndex from pandas.core.indexes.period import PeriodIndex from pandas.core.series import TimedeltaSeries +from typing_extensions import Self from pandas._libs import ( Timedelta, @@ -31,10 +31,7 @@ from pandas._typing import ( num, ) -# type ignore needed because of __getitem__() -class TimedeltaIndex( # type: ignore[misc] - _IndexGetitemMixin[Timedelta], DatetimeTimedeltaMixin, TimedeltaIndexProperties -): +class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties): def __init__( self, data: AnyArrayLike @@ -54,11 +51,11 @@ class TimedeltaIndex( # type: ignore[misc] @overload def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ... @overload - def __add__(self, other: Timedelta | TimedeltaIndex) -> TimedeltaIndex: ... + def __add__(self, other: Timedelta | Self) -> Self: ... def __radd__(self, other: Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override] - def __sub__(self, other: Timedelta | TimedeltaIndex) -> TimedeltaIndex: ... - def __mul__(self, other: num) -> TimedeltaIndex: ... - def __truediv__(self, other: num) -> TimedeltaIndex: ... + def __sub__(self, other: Timedelta | Self) -> Self: ... + def __mul__(self, other: num) -> Self: ... + def __truediv__(self, other: num) -> Self: ... def astype(self, dtype, copy: bool = ...): ... def get_value(self, series, key): ... def get_loc(self, key, tolerance=...): ... diff --git a/pandas-stubs/core/reshape/tile.pyi b/pandas-stubs/core/reshape/tile.pyi index 7d8e16b5a..e561874c7 100644 --- a/pandas-stubs/core/reshape/tile.pyi +++ b/pandas-stubs/core/reshape/tile.pyi @@ -15,13 +15,10 @@ from pandas import ( Series, Timestamp, ) -from pandas.core.indexes.base import ( - _FloatIndexType, - _IntIndexType, -) from pandas.core.series import TimestampSeries from pandas._typing import ( + IntervalT, Label, npt, ) @@ -29,12 +26,7 @@ from pandas._typing import ( @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: int - | Series - | _IntIndexType - | _FloatIndexType - | Sequence[int] - | Sequence[float], + bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], right: bool = ..., *, labels: Literal[False], @@ -47,7 +39,7 @@ def cut( @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: IntervalIndex, + bins: IntervalIndex[IntervalT], right: bool = ..., *, labels: Literal[False], @@ -56,7 +48,7 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[npt.NDArray[np.intp], IntervalIndex]: ... +) -> tuple[npt.NDArray[np.intp], IntervalIndex[IntervalT]]: ... @overload def cut( # type: ignore[misc] x: TimestampSeries, @@ -90,12 +82,7 @@ def cut( @overload def cut( x: Series, - bins: int - | Series - | _IntIndexType - | _FloatIndexType - | Sequence[int] - | Sequence[float], + bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], right: bool = ..., labels: Literal[False] | Sequence[Label] | None = ..., *, @@ -121,12 +108,7 @@ def cut( @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: int - | Series - | _IntIndexType - | _FloatIndexType - | Sequence[int] - | Sequence[float], + bins: int | Series | Index[int] | Index[float] | Sequence[int] | Sequence[float], right: bool = ..., labels: Sequence[Label] | None = ..., *, @@ -139,7 +121,7 @@ def cut( @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - bins: IntervalIndex, + bins: IntervalIndex[IntervalT], right: bool = ..., labels: Sequence[Label] | None = ..., *, @@ -148,14 +130,14 @@ def cut( include_lowest: bool = ..., duplicates: Literal["raise", "drop"] = ..., ordered: bool = ..., -) -> tuple[Categorical, IntervalIndex]: ... +) -> tuple[Categorical, IntervalIndex[IntervalT]]: ... @overload def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], bins: int | Series - | _IntIndexType - | _FloatIndexType + | Index[int] + | Index[float] | Sequence[int] | Sequence[float] | IntervalIndex, @@ -190,8 +172,8 @@ def cut( x: Series, bins: int | Series - | _IntIndexType - | _FloatIndexType + | Index[int] + | Index[float] | Sequence[int] | Sequence[float] | IntervalIndex, @@ -208,8 +190,8 @@ def cut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], bins: int | Series - | _IntIndexType - | _FloatIndexType + | Index[int] + | Index[float] | Sequence[int] | Sequence[float] | IntervalIndex, @@ -224,7 +206,7 @@ def cut( @overload def qcut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, *, labels: Literal[False], retbins: Literal[False] = ..., @@ -234,7 +216,7 @@ def qcut( @overload def qcut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, labels: Sequence[Label] | None = ..., retbins: Literal[False] = ..., precision: int = ..., @@ -243,7 +225,7 @@ def qcut( @overload def qcut( x: Series, - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, labels: Literal[False] | Sequence[Label] | None = ..., retbins: Literal[False] = ..., precision: int = ..., @@ -252,7 +234,7 @@ def qcut( @overload def qcut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, *, labels: Literal[False], retbins: Literal[True], @@ -262,7 +244,7 @@ def qcut( @overload def qcut( x: Series, - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, labels: Literal[False] | Sequence[Label] | None = ..., *, retbins: Literal[True], @@ -272,7 +254,7 @@ def qcut( @overload def qcut( x: Index | npt.NDArray | Sequence[int] | Sequence[float], - q: int | Sequence[float] | Series[float] | _FloatIndexType | npt.NDArray, + q: int | Sequence[float] | Series[float] | Index[float] | npt.NDArray, labels: Sequence[Label] | None = ..., *, retbins: Literal[True], diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 813491e0f..df9db7f42 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -211,7 +211,7 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): value: S1 | ArrayLike | Series[S1] | None, ) -> None: ... -class Series(IndexOpsMixin, NDFrame, Generic[S1]): +class Series(IndexOpsMixin[S1], NDFrame): _ListLike: TypeAlias = ArrayLike | dict[_str, np.ndarray] | list | tuple | Index __hash__: ClassVar[None] diff --git a/tests/test_frame.py b/tests/test_frame.py index 51d901fdd..91263bd7b 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1706,7 +1706,7 @@ def test_indexslice_getitem(): .set_index(["x", "y"]) ) ind = pd.Index([2, 3]) - check(assert_type(pd.IndexSlice[ind, :], "tuple[pd.Index, slice]"), tuple) + check(assert_type(pd.IndexSlice[ind, :], "tuple[pd.Index[int], slice]"), tuple) check(assert_type(df.loc[pd.IndexSlice[ind, :]], pd.DataFrame), pd.DataFrame) check(assert_type(df.loc[pd.IndexSlice[1:2]], pd.DataFrame), pd.DataFrame) check( diff --git a/tests/test_indexes.py b/tests/test_indexes.py index a089e9677..80ec5b3be 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -35,17 +35,6 @@ SupportsRichComparison: TypeAlias = Any else: from _typeshed import SupportsRichComparison # noqa: F401 - from pandas.core.indexes.base import ( - _ComplexIndexType, - _FloatIndexType, - _IntIndexType, - ) -else: - from pandas.core.indexes.base import ( - Index as _ComplexIndexType, - Index as _FloatIndexType, - Index as _IntIndexType, - ) def test_index_unique() -> None: @@ -91,8 +80,8 @@ def test_multiindex_get_level_values() -> None: def test_index_tolist() -> None: i1 = pd.Index([1, 2, 3]) - check(assert_type(i1.tolist(), list), list, int) - check(assert_type(i1.to_list(), list), list, int) + check(assert_type(i1.tolist(), list[int]), list, int) + check(assert_type(i1.to_list(), list[int]), list, int) def test_column_getitem() -> None: @@ -128,23 +117,23 @@ def test_column_sequence() -> None: def test_difference_none() -> None: # https://github.com/pandas-dev/pandas-stubs/issues/17 ind = pd.Index([1, 2, 3]) - check(assert_type(ind.difference([1, None]), pd.Index), pd.Index) + check(assert_type(ind.difference([1, None]), "pd.Index[int]"), pd.Index) # GH 253 - check(assert_type(ind.difference([1]), pd.Index), pd.Index) + check(assert_type(ind.difference([1]), "pd.Index[int]"), pd.Index) def test_str_split() -> None: # GH 194 ind = pd.Index(["a-b", "c-d"]) - check(assert_type(ind.str.split("-"), pd.Index), pd.Index) + check(assert_type(ind.str.split("-"), "pd.Index[str]"), pd.Index) check(assert_type(ind.str.split("-", expand=True), pd.MultiIndex), pd.MultiIndex) def test_index_dropna(): idx = pd.Index([1, 2]) - check(assert_type(idx.dropna(how="all"), pd.Index), pd.Index) - check(assert_type(idx.dropna(how="any"), pd.Index), pd.Index) + check(assert_type(idx.dropna(how="all"), "pd.Index[int]"), pd.Index) + check(assert_type(idx.dropna(how="any"), "pd.Index[int]"), pd.Index) midx = pd.MultiIndex.from_arrays([[1, 2], [3, 4]]) @@ -155,7 +144,7 @@ def test_index_dropna(): def test_index_neg(): # GH 253 idx = pd.Index([1, 2]) - check(assert_type(-idx, pd.Index), pd.Index) + check(assert_type(-idx, "pd.Index[int]"), pd.Index) def test_types_to_numpy() -> None: @@ -168,16 +157,16 @@ def test_types_to_numpy() -> None: def test_index_arithmetic() -> None: # GH 287 idx = pd.Index([1, 2.2, 3], dtype=float) - check(assert_type(idx + 3, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(idx - 3, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(idx * 3, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(idx / 3, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(idx // 3, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(3 + idx, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(3 - idx, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(3 * idx, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(3 / idx, "_FloatIndexType"), _FloatIndexType, np.float64) - check(assert_type(3 // idx, "_FloatIndexType"), _FloatIndexType, np.float64) + check(assert_type(idx + 3, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(idx - 3, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(idx * 3, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(idx / 3, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(idx // 3, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(3 + idx, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(3 - idx, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(3 * idx, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(3 / idx, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(3 // idx, "pd.Index[float]"), pd.Index[float], np.float64) def test_index_relops() -> None: @@ -215,21 +204,21 @@ def test_range_index_union(): check( assert_type( pd.RangeIndex(0, 10).union(pd.RangeIndex(10, 20)), - Union[pd.Index, _IntIndexType, pd.RangeIndex], + Union[pd.Index, pd.Index[int], pd.RangeIndex], ), pd.RangeIndex, ) check( assert_type( pd.RangeIndex(0, 10).union([11, 12, 13]), - Union[pd.Index, _IntIndexType, pd.RangeIndex], + Union[pd.Index, pd.Index[int], pd.RangeIndex], ), pd.Index, ) check( assert_type( pd.RangeIndex(0, 10).union(["a", "b", "c"]), - Union[pd.Index, _IntIndexType, pd.RangeIndex], + Union[pd.Index, pd.Index[int], pd.RangeIndex], ), pd.Index, ) @@ -753,30 +742,30 @@ def test_index_operators() -> None: i1 = pd.Index([1, 2, 3]) i2 = pd.Index([4, 5, 6]) - check(assert_type(i1 + i2, pd.Index), pd.Index) - check(assert_type(i1 + 10, pd.Index), pd.Index) - check(assert_type(10 + i1, pd.Index), pd.Index) - check(assert_type(i1 - i2, pd.Index), pd.Index) - check(assert_type(i1 - 10, pd.Index), pd.Index) - check(assert_type(10 - i1, pd.Index), pd.Index) - check(assert_type(i1 * i2, pd.Index), pd.Index) - check(assert_type(i1 * 10, pd.Index), pd.Index) - check(assert_type(10 * i1, pd.Index), pd.Index) - check(assert_type(i1 / i2, pd.Index), pd.Index) - check(assert_type(i1 / 10, pd.Index), pd.Index) - check(assert_type(10 / i1, pd.Index), pd.Index) - check(assert_type(i1 // i2, pd.Index), pd.Index) - check(assert_type(i1 // 10, pd.Index), pd.Index) - check(assert_type(10 // i1, pd.Index), pd.Index) - check(assert_type(i1**i2, pd.Index), pd.Index) - check(assert_type(i1**2, pd.Index), pd.Index) - check(assert_type(2**i1, pd.Index), pd.Index) - check(assert_type(i1 % i2, pd.Index), pd.Index) - check(assert_type(i1 % 10, pd.Index), pd.Index) - check(assert_type(10 % i1, pd.Index), pd.Index) - check(assert_type(divmod(i1, i2), Tuple[pd.Index, pd.Index]), tuple) - check(assert_type(divmod(i1, 10), Tuple[pd.Index, pd.Index]), tuple) - check(assert_type(divmod(10, i1), Tuple[pd.Index, pd.Index]), tuple) + check(assert_type(i1 + i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 + 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 + i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1 - i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 - 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 - i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1 * i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 * 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 * i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1 / i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 / 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 / i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1 // i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 // 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 // i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1**i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1**2, "pd.Index[int]"), pd.Index) + check(assert_type(2**i1, "pd.Index[int]"), pd.Index) + check(assert_type(i1 % i2, "pd.Index[int]"), pd.Index) + check(assert_type(i1 % 10, "pd.Index[int]"), pd.Index) + check(assert_type(10 % i1, "pd.Index[int]"), pd.Index) + check(assert_type(divmod(i1, i2), Tuple["pd.Index[int]", "pd.Index[int]"]), tuple) + check(assert_type(divmod(i1, 10), Tuple["pd.Index[int]", "pd.Index[int]"]), tuple) + check(assert_type(divmod(10, i1), Tuple["pd.Index[int]", "pd.Index[int]"]), tuple) if TYPE_CHECKING_INVALID_USAGE: assert_type( @@ -905,9 +894,9 @@ def test_getitem() -> None: check(assert_type(mi[[0, 2]], pd.MultiIndex), pd.MultiIndex, tuple) i0 = pd.Index(["a", "b", "c"]) - check(assert_type(i0, pd.Index), pd.Index) - check(assert_type(i0[0], Scalar), str) - check(assert_type(i0[[0, 2]], pd.Index), pd.Index, str) + check(assert_type(i0, "pd.Index[str]"), pd.Index) + check(assert_type(i0[0], str), str) + check(assert_type(i0[[0, 2]], "pd.Index[str]"), pd.Index, str) def test_multiindex_dtypes(): @@ -921,66 +910,80 @@ def test_index_constructors(): # Eventually should be using a generic index ilist = [1, 2, 3] check( - assert_type(pd.Index(ilist, dtype="int"), _IntIndexType), pd.Index, np.integer + assert_type(pd.Index(ilist, dtype="int"), "pd.Index[int]"), pd.Index, np.integer + ) + check( + assert_type(pd.Index(ilist, dtype=int), "pd.Index[int]"), pd.Index, np.integer + ) + check( + assert_type(pd.Index(ilist, dtype=np.int8), "pd.Index[int]"), pd.Index, np.int8 ) - check(assert_type(pd.Index(ilist, dtype=int), _IntIndexType), pd.Index, np.integer) - check(assert_type(pd.Index(ilist, dtype=np.int8), _IntIndexType), pd.Index, np.int8) check( - assert_type(pd.Index(ilist, dtype=np.int16), _IntIndexType), pd.Index, np.int16 + assert_type(pd.Index(ilist, dtype=np.int16), "pd.Index[int]"), + pd.Index, + np.int16, ) check( - assert_type(pd.Index(ilist, dtype=np.int32), _IntIndexType), pd.Index, np.int32 + assert_type(pd.Index(ilist, dtype=np.int32), "pd.Index[int]"), + pd.Index, + np.int32, ) check( - assert_type(pd.Index(ilist, dtype=np.int64), _IntIndexType), pd.Index, np.int64 + assert_type(pd.Index(ilist, dtype=np.int64), "pd.Index[int]"), + pd.Index, + np.int64, ) check( - assert_type(pd.Index(ilist, dtype=np.uint8), _IntIndexType), pd.Index, np.uint8 + assert_type(pd.Index(ilist, dtype=np.uint8), "pd.Index[int]"), + pd.Index, + np.uint8, ) check( - assert_type(pd.Index(ilist, dtype=np.uint16), _IntIndexType), + assert_type(pd.Index(ilist, dtype=np.uint16), "pd.Index[int]"), pd.Index, np.uint16, ) check( - assert_type(pd.Index(ilist, dtype=np.uint32), _IntIndexType), + assert_type(pd.Index(ilist, dtype=np.uint32), "pd.Index[int]"), pd.Index, np.uint32, ) check( - assert_type(pd.Index(ilist, dtype=np.uint64), _IntIndexType), + assert_type(pd.Index(ilist, dtype=np.uint64), "pd.Index[int]"), pd.Index, np.uint64, ) flist = [1.1, 2.2, 3.3] check( - assert_type(pd.Index(flist, dtype="float"), _FloatIndexType), + assert_type(pd.Index(flist, dtype="float"), "pd.Index[float]"), pd.Index, np.float64, ) check( - assert_type(pd.Index(flist, dtype=float), _FloatIndexType), pd.Index, np.float64 + assert_type(pd.Index(flist, dtype=float), "pd.Index[float]"), + pd.Index, + np.float64, ) check( - assert_type(pd.Index(flist, dtype=np.float32), _FloatIndexType), + assert_type(pd.Index(flist, dtype=np.float32), "pd.Index[float]"), pd.Index, np.float32, ) check( - assert_type(pd.Index(flist, dtype=np.float64), _FloatIndexType), + assert_type(pd.Index(flist, dtype=np.float64), "pd.Index[float]"), pd.Index, np.float64, ) clist = [1 + 1j, 2 + 2j, 3 + 4j] check( - assert_type(pd.Index(clist, dtype="complex"), _ComplexIndexType), + assert_type(pd.Index(clist, dtype="complex"), "pd.Index[complex]"), pd.Index, complex, ) check( - assert_type(pd.Index(clist, dtype=complex), _ComplexIndexType), + assert_type(pd.Index(clist, dtype=complex), "pd.Index[complex]"), pd.Index, complex, ) diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 401fc1aa1..bc194bb8a 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -899,7 +899,7 @@ def test_cut() -> None: n0, n1 = pd.cut([1, 2, 3, 4, 5, 6, 7, 8], intval_idx, retbins=True) check(assert_type(n0, pd.Categorical), pd.Categorical) - check(assert_type(n1, pd.IntervalIndex), pd.IntervalIndex) + check(assert_type(n1, "pd.IntervalIndex[pd.Interval[int]]"), pd.IntervalIndex) s1 = pd.Series(data=pd.date_range("1/1/2020", periods=300)) check( diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 6b045950c..adefb2100 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -45,11 +45,6 @@ from pandas.core.series import TimedeltaSeries # noqa: F401 from pandas.core.series import TimestampSeries # noqa: F401 -if TYPE_CHECKING: - from pandas.core.indexes.base import _IntIndexType - -else: - from pandas import Index as _IntIndexType # Separately define here so pytest works np_ndarray_bool = npt.NDArray[np.bool_] @@ -448,7 +443,7 @@ def test_series_dt_accessors() -> None: def test_datetimeindex_accessors() -> None: # GH 194 x = pd.DatetimeIndex(["2022-08-14", "2022-08-20"]) - check(assert_type(x.month, "_IntIndexType"), _IntIndexType) + check(assert_type(x.month, "pd.Index[int]"), pd.Index) i0 = pd.date_range(start="2022-06-01", periods=10) check(assert_type(i0, pd.DatetimeIndex), pd.DatetimeIndex, pd.Timestamp) @@ -456,20 +451,20 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.date, np.ndarray), np.ndarray, dt.date) check(assert_type(i0.time, np.ndarray), np.ndarray, dt.time) check(assert_type(i0.timetz, np.ndarray), np.ndarray, dt.time) - check(assert_type(i0.year, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.month, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.day, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.hour, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.minute, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.second, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.microsecond, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.nanosecond, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.dayofweek, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.day_of_week, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.weekday, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.dayofyear, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.day_of_year, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.quarter, "_IntIndexType"), _IntIndexType, np.int32) + check(assert_type(i0.year, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.month, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.day, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.hour, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.minute, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.second, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.microsecond, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.nanosecond, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.dayofweek, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.day_of_week, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.weekday, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.dayofyear, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.day_of_year, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.quarter, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.is_month_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_month_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_quarter_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) @@ -477,8 +472,8 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.is_year_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_year_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_leap_year, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.daysinmonth, "_IntIndexType"), _IntIndexType, np.int32) - check(assert_type(i0.days_in_month, "_IntIndexType"), _IntIndexType, np.int32) + check(assert_type(i0.daysinmonth, pd.Index[int]), pd.Index, np.int32) + check(assert_type(i0.days_in_month, pd.Index[int]), pd.Index, np.int32) check(assert_type(i0.tz, Optional[dt.tzinfo]), type(None)) check(assert_type(i0.freq, Optional[BaseOffset]), BaseOffset) check(assert_type(i0.isocalendar(), pd.DataFrame), pd.DataFrame) @@ -538,20 +533,20 @@ def test_periodindex_accessors() -> None: i0 = pd.period_range(start="2022-06-01", periods=10) check(assert_type(i0, pd.PeriodIndex), pd.PeriodIndex, pd.Period) - check(assert_type(i0.year, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.month, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.day, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.hour, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.minute, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.second, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.dayofweek, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.day_of_week, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.weekday, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.dayofyear, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.day_of_year, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.quarter, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.daysinmonth, "_IntIndexType"), _IntIndexType, np.integer) - check(assert_type(i0.days_in_month, "_IntIndexType"), _IntIndexType, np.integer) + check(assert_type(i0.year, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.month, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.day, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.hour, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.minute, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.second, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.dayofweek, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.day_of_week, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.weekday, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.dayofyear, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.day_of_year, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.quarter, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.daysinmonth, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.days_in_month, pd.Index[int]), pd.Index, np.integer) check(assert_type(i0.freq, Optional[BaseOffset]), BaseOffset) check(assert_type(i0.strftime("%Y"), pd.Index), pd.Index, str) check(assert_type(i0.asfreq("D"), pd.PeriodIndex), pd.PeriodIndex, pd.Period) From 9e40ddfdd040e858ae08a1b7adf0b30b7516033a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 10:39:43 -0400 Subject: [PATCH 03/23] mypy handles Never differently - at least asssert that the method never returns --- tests/test_scalars.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 8c946839d..1ef0e543f 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -15,6 +15,7 @@ import pandas as pd import pytz from typing_extensions import ( + Never, TypeAlias, assert_type, ) @@ -745,8 +746,12 @@ def test_timedelta_mul_div() -> None: md_ndarray_float // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] mp_series_int // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - md_int64_index // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - md_float_index // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + assert_type( + md_int64_index // td, Never # pyright: ignore[reportGeneralTypeIssues] + ) + assert_type( + md_float_index // td, Never # pyright: ignore[reportGeneralTypeIssues] + ) check(assert_type(td / td, float), float) check(assert_type(td / pd.NaT, float), float) @@ -779,8 +784,12 @@ def test_timedelta_mul_div() -> None: md_ndarray_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] mp_series_int / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - md_int64_index / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - md_float_index / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + assert_type( + md_int64_index / td, Never # pyright: ignore[reportGeneralTypeIssues] + ) + assert_type( + md_float_index / td, Never # pyright: ignore[reportGeneralTypeIssues] + ) def test_timedelta_mod_abs_unary() -> None: From 68abe4da241c22fba1c0eae99c77e134cb10b5f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 13:52:05 -0400 Subject: [PATCH 04/23] use Self --- pandas-stubs/core/indexes/base.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 106009cd9..1f53b8b3c 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -107,7 +107,7 @@ class Index(IndexOpsMixin[S1], PandasObject): name=..., tupleize_cols: bool = ..., **kwargs, - ) -> Index: ... + ) -> Self: ... @property def str(self) -> StringMethods[Self, MultiIndex]: ... @property From 2cfbdb0fd9cae7aab95c26435638502fafa8f2a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 14:30:04 -0400 Subject: [PATCH 05/23] a few more S1/Self --- pandas-stubs/core/indexes/base.pyi | 13 ++++++------- tests/test_indexes.py | 25 ++----------------------- 2 files changed, 8 insertions(+), 30 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 1f53b8b3c..eff99f923 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -42,7 +42,6 @@ from pandas._typing import ( Label, Level, NaPosition, - Scalar, np_ndarray_anyint, np_ndarray_bool, np_ndarray_int64, @@ -193,7 +192,7 @@ class Index(IndexOpsMixin[S1], PandasObject): def __nonzero__(self) -> None: ... __bool__ = ... def union(self, other: list[HashableT] | Index, sort=...) -> Index: ... - def intersection(self, other: list[T1] | Self, sort: bool = ...) -> Self: ... + def intersection(self, other: list[T1] | Index, sort: bool = ...) -> Self: ... def difference(self, other: list | Index, sort: bool | None = None) -> Self: ... def symmetric_difference( self, other: list[T1] | Index, result_name=..., sort=... @@ -264,12 +263,12 @@ class Index(IndexOpsMixin[S1], PandasObject): def shape(self) -> tuple[int, ...]: ... # Extra methods from old stubs def __eq__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] - def __iter__(self) -> Iterator: ... + def __iter__(self) -> Iterator[S1]: ... def __ne__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] - def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] - def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] - def __lt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] - def __gt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] + def __le__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] + def __ge__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] + def __lt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] + def __gt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] # overwrite inherit methods from OpsMixin @overload def __mul__( # type: ignore[misc] diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 80ec5b3be..90b9bcbee 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -25,16 +25,6 @@ if TYPE_CHECKING: MYPY_CHECKING: bool = True - # See test_sorted_and_list() where mypy and pyright do different - # inference on sorted(pd.Index) - if MYPY_CHECKING: - from typing import Any - - from typing_extensions import TypeAlias - - SupportsRichComparison: TypeAlias = Any - else: - from _typeshed import SupportsRichComparison # noqa: F401 def test_index_unique() -> None: @@ -719,22 +709,11 @@ def test_interval_index_tuples(): def test_sorted_and_list() -> None: # GH 497 i1 = pd.Index([3, 2, 1]) - # mypy infers sorted(i1) as list[Any], while pyright infers sorted(i1) as - # list[SupportsRichComparison] check( - assert_type( - sorted(i1), - "list[SupportsRichComparison]", - ), - list, - ) - check( - assert_type( - list(i1), - list, - ), + assert_type(sorted(i1), list[int]), list, ) + check(assert_type(list(i1), list[int]), list) def test_index_operators() -> None: From ebcc389e6ef5d5c671d27d11b7854e47509a23ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 15:04:24 -0400 Subject: [PATCH 06/23] fix tests --- pandas-stubs/core/indexes/base.pyi | 5 ++--- tests/test_indexes.py | 34 ++++++++++++++++++++++-------- tests/test_timefuncs.py | 32 ++++++++++++++-------------- 3 files changed, 43 insertions(+), 28 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index eff99f923..54feb5a3e 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -33,7 +33,6 @@ from typing_extensions import ( from pandas._typing import ( S1, - T1, Dtype, DtypeArg, DtypeObj, @@ -192,10 +191,10 @@ class Index(IndexOpsMixin[S1], PandasObject): def __nonzero__(self) -> None: ... __bool__ = ... def union(self, other: list[HashableT] | Index, sort=...) -> Index: ... - def intersection(self, other: list[T1] | Index, sort: bool = ...) -> Self: ... + def intersection(self, other: list[S1] | Self, sort: bool = ...) -> Self: ... def difference(self, other: list | Index, sort: bool | None = None) -> Self: ... def symmetric_difference( - self, other: list[T1] | Index, result_name=..., sort=... + self, other: list[S1] | Self, result_name=..., sort=... ) -> Self: ... def get_loc( self, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 90b9bcbee..fa4656497 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -151,12 +151,12 @@ def test_index_arithmetic() -> None: check(assert_type(idx - 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx * 3, "pd.Index[float]"), pd.Index, np.float64) check(assert_type(idx / 3, "pd.Index[float]"), pd.Index, np.float64) - check(assert_type(idx // 3, "pd.Index[float]"), pd.Index[float], np.float64) - check(assert_type(3 + idx, "pd.Index[float]"), pd.Index[float], np.float64) - check(assert_type(3 - idx, "pd.Index[float]"), pd.Index[float], np.float64) - check(assert_type(3 * idx, "pd.Index[float]"), pd.Index[float], np.float64) - check(assert_type(3 / idx, "pd.Index[float]"), pd.Index[float], np.float64) - check(assert_type(3 // idx, "pd.Index[float]"), pd.Index[float], np.float64) + check(assert_type(idx // 3, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(3 + idx, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(3 - idx, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(3 * idx, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(3 / idx, "pd.Index[float]"), pd.Index, np.float64) + check(assert_type(3 // idx, "pd.Index[float]"), pd.Index, np.float64) def test_index_relops() -> None: @@ -194,21 +194,21 @@ def test_range_index_union(): check( assert_type( pd.RangeIndex(0, 10).union(pd.RangeIndex(10, 20)), - Union[pd.Index, pd.Index[int], pd.RangeIndex], + Union[pd.Index, "pd.Index[int]", pd.RangeIndex], ), pd.RangeIndex, ) check( assert_type( pd.RangeIndex(0, 10).union([11, 12, 13]), - Union[pd.Index, pd.Index[int], pd.RangeIndex], + Union[pd.Index, "pd.Index[int]", pd.RangeIndex], ), pd.Index, ) check( assert_type( pd.RangeIndex(0, 10).union(["a", "b", "c"]), - Union[pd.Index, pd.Index[int], pd.RangeIndex], + Union[pd.Index, "pd.Index[int]", pd.RangeIndex], ), pd.Index, ) @@ -973,3 +973,19 @@ def test_index_constructors(): # to specify all the possible dtype options. For right now, we will leave the # test here as a reminder that we would like this to be seen as incorrect usage. pd.Index(flist, dtype=np.float16) + + +def test_iter() -> None: + # GH 723 + for ts in pd.date_range(start="1/1/2023", end="1/08/2023", freq="6H"): + check(assert_type(ts, pd.Timestamp), pd.Timestamp) + + +def test_intersection() -> None: + # GH 744 + index = pd.DatetimeIndex(["2022-01-01"]) + check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex) + check( + assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex), + pd.DatetimeIndex, + ) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index adefb2100..4977d43ca 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -472,8 +472,8 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.is_year_start, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_year_end, npt.NDArray[np.bool_]), np.ndarray, np.bool_) check(assert_type(i0.is_leap_year, npt.NDArray[np.bool_]), np.ndarray, np.bool_) - check(assert_type(i0.daysinmonth, pd.Index[int]), pd.Index, np.int32) - check(assert_type(i0.days_in_month, pd.Index[int]), pd.Index, np.int32) + check(assert_type(i0.daysinmonth, "pd.Index[int]"), pd.Index, np.int32) + check(assert_type(i0.days_in_month, "pd.Index[int]"), pd.Index, np.int32) check(assert_type(i0.tz, Optional[dt.tzinfo]), type(None)) check(assert_type(i0.freq, Optional[BaseOffset]), BaseOffset) check(assert_type(i0.isocalendar(), pd.DataFrame), pd.DataFrame) @@ -533,20 +533,20 @@ def test_periodindex_accessors() -> None: i0 = pd.period_range(start="2022-06-01", periods=10) check(assert_type(i0, pd.PeriodIndex), pd.PeriodIndex, pd.Period) - check(assert_type(i0.year, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.month, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.day, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.hour, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.minute, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.second, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.dayofweek, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.day_of_week, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.weekday, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.dayofyear, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.day_of_year, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.quarter, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.daysinmonth, pd.Index[int]), pd.Index, np.integer) - check(assert_type(i0.days_in_month, pd.Index[int]), pd.Index, np.integer) + check(assert_type(i0.year, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.month, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.day, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.hour, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.minute, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.second, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.dayofweek, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.day_of_week, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.weekday, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.dayofyear, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.day_of_year, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.quarter, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.daysinmonth, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(i0.days_in_month, "pd.Index[int]"), pd.Index, np.integer) check(assert_type(i0.freq, Optional[BaseOffset]), BaseOffset) check(assert_type(i0.strftime("%Y"), pd.Index), pd.Index, str) check(assert_type(i0.asfreq("D"), pd.PeriodIndex), pd.PeriodIndex, pd.Period) From 24d42fa505054d0359797496215ad6a13f866aaa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 15:47:07 -0400 Subject: [PATCH 07/23] more tests --- tests/test_frame.py | 3 ++- tests/test_indexes.py | 20 ++++++++++++++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 91263bd7b..46a689aa2 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1761,7 +1761,8 @@ def test_getmultiindex_columns() -> None: [(i, s) for i in [1] for s in df.columns.get_level_values(1)] ] res4: pd.DataFrame = df[[df.columns[0]]] - check(assert_type(df[df.columns[0]], pd.Series), pd.Series) + column: Scalar = df.columns[0] + check(assert_type(df[column], pd.Series), pd.Series) check(assert_type(df[li[0]], pd.Series), pd.Series) diff --git a/tests/test_indexes.py b/tests/test_indexes.py index fa4656497..a292604ee 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -16,7 +16,6 @@ ) from pandas._typing import Dtype # noqa: F401 -from pandas._typing import Scalar from tests import ( TYPE_CHECKING_INVALID_USAGE, @@ -78,9 +77,9 @@ def test_column_getitem() -> None: # https://github.com/microsoft/python-type-stubs/issues/199#issuecomment-1132806594 df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"]) - column = df.columns[0] - check(assert_type(column, Scalar), str) - check(assert_type(df[column], pd.Series), pd.Series, np.int64) + columns: pd.Index[str] = df.columns + check(assert_type(columns[0], str), str) + check(assert_type(df[columns[0]], pd.Series), pd.Series, np.int64) def test_column_contains() -> None: @@ -989,3 +988,16 @@ def test_intersection() -> None: assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex), pd.DatetimeIndex, ) + + +def test_annotate() -> None: + # GH 502 + df = pd.DataFrame({"a": [1, 2]}) + + columns: pd.Index[str] = df.columns + for column in columns: + check(assert_type(column, str), str) + + names: list[str] = list(df.columns) + for name in names: + check(assert_type(name, str), str) From f3802d3cd11b56c935a1030a3de45a7a591c9c2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 16:09:51 -0400 Subject: [PATCH 08/23] add an ignore to let mypy pass --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 7f5652073..457c587f6 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -251,8 +251,9 @@ class Timedelta(timedelta): def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ... @overload def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ... + # maybe related to https://github.com/python/mypy/issues/10755 @overload - def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... + def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ... # type: ignore[misc] # Override due to more types supported than dt.timedelta # error: Signature of "__floordiv__" incompatible with supertype "timedelta" @overload # type: ignore[override] From 92eab418065ad66ef5fd34139b79f0d3d73a9f88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 16:14:51 -0400 Subject: [PATCH 09/23] compat with old python versions --- tests/test_indexes.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/tests/test_indexes.py b/tests/test_indexes.py index a292604ee..e84ee2ecc 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -69,8 +69,8 @@ def test_multiindex_get_level_values() -> None: def test_index_tolist() -> None: i1 = pd.Index([1, 2, 3]) - check(assert_type(i1.tolist(), list[int]), list, int) - check(assert_type(i1.to_list(), list[int]), list, int) + check(assert_type(i1.tolist(), "list[int]"), list, int) + check(assert_type(i1.to_list(), "list[int]"), list, int) def test_column_getitem() -> None: @@ -708,11 +708,8 @@ def test_interval_index_tuples(): def test_sorted_and_list() -> None: # GH 497 i1 = pd.Index([3, 2, 1]) - check( - assert_type(sorted(i1), list[int]), - list, - ) - check(assert_type(list(i1), list[int]), list) + check(assert_type(sorted(i1), "list[int]"), list, int) + check(assert_type(list(i1), "list[int]"), list, int) def test_index_operators() -> None: From 6e04d8e4c5c3b04f029b866ac85d663b15075f4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 16:44:13 -0400 Subject: [PATCH 10/23] fix CategoricalIndex; MultiIndex should probably be Index[tuple[S1, ...]] but keeping it as Index[Any] --- pandas-stubs/core/indexes/category.pyi | 12 ++++++++---- pandas-stubs/core/indexes/multi.pyi | 21 +++++++++++---------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/pandas-stubs/core/indexes/category.pyi b/pandas-stubs/core/indexes/category.pyi index b2bcfa615..59dcbf185 100644 --- a/pandas-stubs/core/indexes/category.pyi +++ b/pandas-stubs/core/indexes/category.pyi @@ -1,19 +1,23 @@ +from collections.abc import Iterable from typing import Literal import numpy as np from pandas.core import accessor -from pandas.core.indexes.base import Index # , maybe_extract_name +from pandas.core.indexes.base import Index from pandas.core.indexes.extension import ExtensionIndex from typing_extensions import Self -from pandas._typing import DtypeArg +from pandas._typing import ( + S1, + DtypeArg, +) -class CategoricalIndex(ExtensionIndex, accessor.PandasDelegate): +class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate): codes: np.ndarray = ... categories: Index = ... def __new__( cls, - data=..., + data: Iterable[S1] = ..., categories=..., ordered=..., dtype=..., diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 219ddf9aa..f3afc402a 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -4,6 +4,7 @@ from collections.abc import ( Sequence, ) from typing import ( + Any, Literal, overload, ) @@ -22,7 +23,7 @@ from pandas._typing import ( np_ndarray_bool, ) -class MultiIndex(Index): +class MultiIndex(Index[Any]): def __new__( cls, levels=..., @@ -48,13 +49,13 @@ class MultiIndex(Index): _set_identity: bool = ..., ) -> None: ... @classmethod - def from_arrays(cls, arrays, sortorder=..., names=...) -> MultiIndex: ... + def from_arrays(cls, arrays, sortorder=..., names=...) -> Self: ... @classmethod - def from_tuples(cls, tuples, sortorder=..., names=...) -> MultiIndex: ... + def from_tuples(cls, tuples, sortorder=..., names=...) -> Self: ... @classmethod - def from_product(cls, iterables, sortorder=..., names=...) -> MultiIndex: ... + def from_product(cls, iterables, sortorder=..., names=...) -> Self: ... @classmethod - def from_frame(cls, df, sortorder=..., names=...) -> MultiIndex: ... + def from_frame(cls, df, sortorder=..., names=...) -> Self: ... @property def shape(self): ... @property # Should be read-only @@ -63,7 +64,7 @@ class MultiIndex(Index): @property def codes(self): ... def set_codes(self, codes, *, level=..., verify_integrity: bool = ...): ... - def copy(self, names=..., deep: bool = ...) -> MultiIndex: ... + def copy(self, names=..., deep: bool = ...) -> Self: ... def __array__(self, dtype=...) -> np.ndarray: ... def view(self, cls=...): ... def __contains__(self, key) -> bool: ... @@ -94,7 +95,7 @@ class MultiIndex(Index): def is_monotonic_decreasing(self) -> bool: ... def duplicated(self, keep: Literal["first", "last", False] = ...): ... def fillna(self, value=..., downcast=...) -> None: ... - def dropna(self, how: Literal["any", "all"] = ...) -> MultiIndex: ... + def dropna(self, how: Literal["any", "all"] = ...) -> Self: ... def get_value(self, series, key): ... def get_level_values(self, level: str | int) -> Index: ... def unique(self, level=...): ... @@ -123,7 +124,7 @@ class MultiIndex(Index): | pd.Series[bool] | Sequence[bool] | np_ndarray_bool, - ) -> MultiIndex: ... + ) -> Self: ... @overload def __getitem__(self, key: int) -> tuple: ... def take( @@ -133,7 +134,7 @@ class MultiIndex(Index): def argsort(self, *args, **kwargs): ... def repeat(self, repeats, axis=...): ... def where(self, cond, other=...) -> None: ... - def drop(self, codes, level=..., errors: str = ...) -> MultiIndex: ... # type: ignore[override] + def drop(self, codes, level=..., errors: str = ...) -> Self: ... # type: ignore[override] def swaplevel(self, i: int = ..., j: int = ...): ... def reorder_levels(self, order): ... def sortlevel( @@ -155,7 +156,7 @@ class MultiIndex(Index): def union(self, other, sort=...): ... def intersection(self, other, sort: bool = ...): ... def difference(self, other, sort=...): ... - def astype(self, dtype: DtypeArg | T1, copy: bool = ...) -> MultiIndex: ... + def astype(self, dtype: DtypeArg | T1, copy: bool = ...) -> Self: ... def insert(self, loc, item): ... def delete(self, loc): ... def isin(self, values, level=...) -> np_ndarray_bool: ... From e7a2d3b69fd03f9cb44f2ef6f999e9c4a42e6477 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 19:32:38 -0400 Subject: [PATCH 11/23] Revert "interpolate method" This reverts commit 4929ecb497c1d66a1ea8489d01bbaaf29dec0f4a. --- .pre-commit-config.yaml | 2 +- pandas-stubs/_typing.pyi | 20 -------------------- pandas-stubs/core/frame.pyi | 7 +++---- pandas-stubs/core/resample.pyi | 27 ++++++++++++++++++++++++--- pandas-stubs/core/series.pyi | 23 +++++++++++++++++++++-- tests/test_scalars.py | 2 +- 6 files changed, 50 insertions(+), 31 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1786d3276..cc0e9aaa0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,7 @@ repos: hooks: - id: isort - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.282 + rev: v0.0.278 hooks: - id: ruff args: [ diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 4306aede0..27548758f 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -664,26 +664,6 @@ StataDateFormat: TypeAlias = Literal[ ] FillnaOptions: TypeAlias = Literal["backfill", "bfill", "ffill", "pad"] -InterpolateOptions: TypeAlias = Literal[ - "linear", - "time", - "index", - "pad", - "nearest", - "zero", - "slinear", - "quadratic", - "cubic", - "barycentric", - "polynomial", - "krogh", - "piecewise_polynomial", - "spline", - "pchip", - "akima", - "cubicspline", - "from_derivatives", -] ReplaceMethod: TypeAlias = Literal["pad", "ffill", "bfill"] SortKind: TypeAlias = Literal["quicksort", "mergesort", "heapsort", "stable"] NaPosition: TypeAlias = Literal["first", "last"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index a8f46facb..f9c195b18 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -86,7 +86,6 @@ from pandas._typing import ( IndexingInt, IndexLabel, IndexType, - InterpolateOptions, IntervalClosedType, IntervalT, JoinHow, @@ -1705,7 +1704,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: InterpolateOptions = ..., + method: _str = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -1718,7 +1717,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: InterpolateOptions = ..., + method: _str = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -1731,7 +1730,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def interpolate( self, - method: InterpolateOptions = ..., + method: _str = ..., *, axis: Axis = ..., limit: int | None = ..., diff --git a/pandas-stubs/core/resample.pyi b/pandas-stubs/core/resample.pyi index a21b7dfcf..23de00bf3 100644 --- a/pandas-stubs/core/resample.pyi +++ b/pandas-stubs/core/resample.pyi @@ -22,7 +22,6 @@ from typing_extensions import TypeAlias from pandas._typing import ( Axis, - InterpolateOptions, NDFrameT, Scalar, npt, @@ -51,6 +50,28 @@ _SeriesGroupByFuncArgs: TypeAlias = ( _SeriesGroupByFuncTypes | Mapping[Hashable, _SeriesGroupByFunc | str] ) +_Interpolation: TypeAlias = Literal[ + "linear", + "time", + "index", + "pad", + "nearest", + "zero", + "slinear", + "quadratic", + "cubic", + "spline", + "barycentric", + "polynomial", + "krogh", + "piecewise_polynomial", + "spline", + "pchip", + "akima", + "cubicspline", + "from_derivatives", +] + class Resampler(BaseGroupBy, Generic[NDFrameT]): def __getattr__(self, attr: str) -> SeriesGroupBy: ... def __iter__(self) -> Generator[tuple[Hashable, NDFrameT], None, None]: ... @@ -126,7 +147,7 @@ class Resampler(BaseGroupBy, Generic[NDFrameT]): @overload def interpolate( self, - method: InterpolateOptions = ..., + method: _Interpolation = ..., *, axis: Axis = ..., limit: int | None = ..., @@ -139,7 +160,7 @@ class Resampler(BaseGroupBy, Generic[NDFrameT]): @overload def interpolate( self, - method: InterpolateOptions = ..., + method: _Interpolation = ..., *, axis: Axis = ..., limit: int | None = ..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index df9db7f42..089dcc43f 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -118,7 +118,6 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntDtypeArg, - InterpolateOptions, IntervalClosedType, IntervalT, JoinHow, @@ -1274,7 +1273,27 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[S1] | None: ... def interpolate( self, - method: InterpolateOptions = ..., + method: _str + | Literal[ + "linear", + "time", + "index", + "values", + "pad", + "nearest", + "slinear", + "quadratic", + "cubic", + "spline", + "barycentric", + "polynomial", + "krogh", + "pecewise_polynomial", + "spline", + "pchip", + "akima", + "from_derivatives", + ] = ..., *, axis: AxisIndex | None = ..., limit: int | None = ..., diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 1ef0e543f..136bc6cb9 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -1754,7 +1754,7 @@ def test_period_add_subtract() -> None: as_dt_td = dt.timedelta(days=1) as_np_td = np.timedelta64(1, "D") as_np_i64 = np.int64(1) - as_int: int = 1 + as_int = int(1) as_period_index = pd.period_range("2012-1-1", periods=10, freq="D") check(assert_type(as_period_index, pd.PeriodIndex), pd.PeriodIndex) as_period = pd.Period("2012-1-1", freq="D") From d12a72a4302dfad65e9d373d981f2437c8cfb2a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 21:25:59 -0400 Subject: [PATCH 12/23] many more overloads for subclasses (I wish pandas would not handle subclasses in parent classes) --- pandas-stubs/core/indexes/base.pyi | 124 ++++++++++++++++++++++++++++- tests/test_indexes.py | 28 +++++++ 2 files changed, 150 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 54feb5a3e..17ea6caad 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -5,7 +5,10 @@ from collections.abc import ( Iterator, Sequence, ) -from datetime import timedelta +from datetime import ( + datetime, + timedelta, +) from typing import ( Any, ClassVar, @@ -16,7 +19,13 @@ from typing import ( import numpy as np from pandas import ( DataFrame, + DatetimeIndex, + Interval, + IntervalIndex, MultiIndex, + Period, + PeriodDtype, + PeriodIndex, Series, TimedeltaIndex, ) @@ -31,6 +40,7 @@ from typing_extensions import ( Self, ) +from pandas._libs.interval import _OrderableT from pandas._typing import ( S1, Dtype, @@ -41,6 +51,8 @@ from pandas._typing import ( Label, Level, NaPosition, + TimedeltaDtypeArg, + TimestampDtypeArg, np_ndarray_anyint, np_ndarray_bool, np_ndarray_int64, @@ -53,10 +65,12 @@ _str = str class Index(IndexOpsMixin[S1], PandasObject): __hash__: ClassVar[None] # type: ignore[assignment] + # overloads with additional dtypes @overload def __new__( # type: ignore[misc] cls, data: Iterable, + *, dtype: Literal["int"] | type_t[int] | type_t[np.integer], copy: bool = ..., name=..., @@ -67,6 +81,7 @@ class Index(IndexOpsMixin[S1], PandasObject): def __new__( # type: ignore[misc] cls, data: Iterable, + *, dtype: Literal["float"] | type_t[float] | type_t[np.float32] @@ -80,16 +95,108 @@ class Index(IndexOpsMixin[S1], PandasObject): def __new__( # type: ignore[misc] cls, data: Iterable, + *, dtype: Literal["complex"] | type_t[complex], copy: bool = ..., name=..., tupleize_cols: bool = ..., **kwargs, ) -> Index[complex]: ... + # special overloads with dedicated Index-subclasses + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[np.datetime64 | datetime], + *, + dtype: TimestampDtypeArg = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> DatetimeIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable, + *, + dtype: TimestampDtypeArg, + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> DatetimeIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[Period], + *, + dtype: PeriodDtype = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> PeriodIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable, + *, + dtype: PeriodDtype, + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> PeriodIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[np.timedelta64 | timedelta], + *, + dtype: TimedeltaDtypeArg = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> TimedeltaIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable, + *, + dtype: TimedeltaDtypeArg, + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> TimedeltaIndex: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[Interval[_OrderableT]], + *, + dtype: Literal["Interval"] = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> IntervalIndex[Interval[_OrderableT]]: ... + @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable, + *, + dtype: Literal["Interval"], + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> IntervalIndex[Interval[Any]]: ... + # generic overloads @overload def __new__( cls, data: Iterable[S1] = ..., + *, dtype=..., copy: bool = ..., name=..., @@ -100,6 +207,19 @@ class Index(IndexOpsMixin[S1], PandasObject): def __new__( cls, data: Iterable = ..., + *, + dtype: type[S1] = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> Self: ... + # fallback overload + @overload + def __new__( + cls, + data: Iterable = ..., + *, dtype=..., copy: bool = ..., name=..., @@ -268,7 +388,7 @@ class Index(IndexOpsMixin[S1], PandasObject): def __ge__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] def __lt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] def __gt__(self, other: Self | S1) -> np_ndarray_bool: ... # type: ignore[override] - # overwrite inherit methods from OpsMixin + # overwrite inherited methods from OpsMixin @overload def __mul__( # type: ignore[misc] self: Index[int] | Index[float], other: timedelta diff --git a/tests/test_indexes.py b/tests/test_indexes.py index e84ee2ecc..627764210 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -998,3 +998,31 @@ def test_annotate() -> None: names: list[str] = list(df.columns) for name in names: check(assert_type(name, str), str) + + +def test_new() -> None: + check(assert_type(pd.Index([1]), "pd.Index[int]"), pd.Index, np.intp) + check(assert_type(pd.Index([1], dtype=float), "pd.Index[float]"), pd.Index, float) + check( + assert_type(pd.Index([pd.Timestamp(0)]), pd.DatetimeIndex), + pd.DatetimeIndex, + pd.Timestamp, + ) + check( + assert_type(pd.Index([pd.Timedelta(0)]), pd.TimedeltaIndex), + pd.TimedeltaIndex, + pd.Timedelta, + ) + check( + assert_type(pd.Index([pd.Period("2012-1-1", freq="D")]), pd.PeriodIndex), + pd.PeriodIndex, + pd.Period, + ) + check( + assert_type( + pd.Index([pd.Interval(pd.Timestamp(0), pd.Timestamp(1))]), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) From 87c62d07a1585b20223495212f80c6dc8179458f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 3 Aug 2023 23:35:06 -0400 Subject: [PATCH 13/23] remove PandasObject - it caused the pyright issues but it provides no methods that aren't already provided by object --- pandas-stubs/core/accessor.pyi | 3 --- pandas-stubs/core/arrays/categorical.pyi | 9 +++------ pandas-stubs/core/arrays/sparse/array.pyi | 3 +-- pandas-stubs/core/base.pyi | 4 ---- pandas-stubs/core/generic.pyi | 3 +-- pandas-stubs/core/groupby/groupby.pyi | 9 +++------ pandas-stubs/core/indexes/accessors.pyi | 7 ++----- pandas-stubs/core/indexes/base.pyi | 7 ++----- pandas-stubs/core/indexes/frozen.pyi | 4 +--- pandas-stubs/io/sql.pyi | 5 ++--- pandas-stubs/plotting/_core.pyi | 3 +-- 11 files changed, 16 insertions(+), 41 deletions(-) diff --git a/pandas-stubs/core/accessor.pyi b/pandas-stubs/core/accessor.pyi index d2b7d78a4..62ff02e3f 100644 --- a/pandas-stubs/core/accessor.pyi +++ b/pandas-stubs/core/accessor.pyi @@ -1,8 +1,5 @@ from typing import Any -class DirNamesMixin: - def __dir__(self): ... - class PandasDelegate: ... def delegate_names( diff --git a/pandas-stubs/core/arrays/categorical.pyi b/pandas-stubs/core/arrays/categorical.pyi index 472d92dce..2d0b16359 100644 --- a/pandas-stubs/core/arrays/categorical.pyi +++ b/pandas-stubs/core/arrays/categorical.pyi @@ -12,10 +12,7 @@ import numpy as np from pandas import Series from pandas.core.accessor import PandasDelegate as PandasDelegate from pandas.core.arrays.base import ExtensionArray as ExtensionArray -from pandas.core.base import ( - NoNewAttributesMixin as NoNewAttributesMixin, - PandasObject as PandasObject, -) +from pandas.core.base import NoNewAttributesMixin as NoNewAttributesMixin from pandas.core.indexes.base import Index from pandas._typing import ( @@ -33,7 +30,7 @@ from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype def contains(cat, key, container): ... -class Categorical(ExtensionArray, PandasObject): +class Categorical(ExtensionArray): __array_priority__: int = ... def __init__( self, @@ -185,7 +182,7 @@ class Categorical(ExtensionArray, PandasObject): def repeat(self, repeats, axis=...): ... def isin(self, values): ... -class CategoricalAccessor(PandasDelegate, PandasObject, NoNewAttributesMixin): +class CategoricalAccessor(PandasDelegate, NoNewAttributesMixin): def __init__(self, data) -> None: ... @property def codes(self) -> Series[int]: ... diff --git a/pandas-stubs/core/arrays/sparse/array.pyi b/pandas-stubs/core/arrays/sparse/array.pyi index b11cfaa4e..2ecf26191 100644 --- a/pandas-stubs/core/arrays/sparse/array.pyi +++ b/pandas-stubs/core/arrays/sparse/array.pyi @@ -3,11 +3,10 @@ from pandas.core.arrays import ( ExtensionArray, ExtensionOpsMixin, ) -from pandas.core.base import PandasObject from pandas._typing import TakeIndexer -class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin): +class SparseArray(ExtensionArray, ExtensionOpsMixin): def __init__( self, data, diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index 865d59e7a..f32c934c4 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -6,7 +6,6 @@ from typing import ( import numpy as np from pandas import Index -from pandas.core.accessor import DirNamesMixin from pandas.core.arraylike import OpsMixin from pandas.core.arrays import ExtensionArray from pandas.core.arrays.categorical import Categorical @@ -21,9 +20,6 @@ from pandas._typing import ( npt, ) -class PandasObject(DirNamesMixin): - def __sizeof__(self) -> int: ... - class NoNewAttributesMixin: def __setattr__(self, key, value) -> None: ... diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 034328e1a..5ff08dff1 100644 --- a/pandas-stubs/core/generic.pyi +++ b/pandas-stubs/core/generic.pyi @@ -16,7 +16,6 @@ from typing import ( import numpy as np from pandas import Index -from pandas.core.base import PandasObject import pandas.core.indexing as indexing from pandas.core.series import Series import sqlalchemy.engine @@ -54,7 +53,7 @@ from pandas.io.sql import SQLTable _bool = bool _str = str -class NDFrame(PandasObject, indexing.IndexingMixin): +class NDFrame(indexing.IndexingMixin): __hash__: ClassVar[None] # type: ignore[assignment] def set_flags( diff --git a/pandas-stubs/core/groupby/groupby.pyi b/pandas-stubs/core/groupby/groupby.pyi index 573f9bfb1..32942ab16 100644 --- a/pandas-stubs/core/groupby/groupby.pyi +++ b/pandas-stubs/core/groupby/groupby.pyi @@ -4,10 +4,7 @@ from collections.abc import ( ) import numpy as np -from pandas.core.base import ( - PandasObject, - SelectionMixin, -) +from pandas.core.base import SelectionMixin from pandas.core.frame import DataFrame from pandas.core.generic import NDFrame from pandas.core.groupby import ops @@ -21,12 +18,12 @@ from pandas._typing import ( npt, ) -class GroupByPlot(PandasObject): +class GroupByPlot: def __init__(self, groupby) -> None: ... def __call__(self, *args, **kwargs): ... def __getattr__(self, name: str): ... -class BaseGroupBy(PandasObject, SelectionMixin[NDFrameT]): +class BaseGroupBy(SelectionMixin[NDFrameT]): level = ... as_index = ... keys = ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index f181f7036..db07d6e99 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -20,10 +20,7 @@ from pandas.core.arrays import ( DatetimeArray, PeriodArray, ) -from pandas.core.base import ( - NoNewAttributesMixin, - PandasObject, -) +from pandas.core.base import NoNewAttributesMixin from pandas.core.frame import DataFrame from pandas.core.series import ( PeriodSeries, @@ -39,7 +36,7 @@ from pandas._typing import ( np_ndarray_bool, ) -class Properties(PandasDelegate, PandasObject, NoNewAttributesMixin): +class Properties(PandasDelegate, NoNewAttributesMixin): def __init__(self, data: Series, orig) -> None: ... _DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", Series[int], Index[int]) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 17ea6caad..2d7c173bc 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -30,10 +30,7 @@ from pandas import ( TimedeltaIndex, ) from pandas.core.arrays import ExtensionArray -from pandas.core.base import ( - IndexOpsMixin, - PandasObject, -) +from pandas.core.base import IndexOpsMixin from pandas.core.strings import StringMethods from typing_extensions import ( Never, @@ -63,7 +60,7 @@ class InvalidIndexError(Exception): ... _str = str -class Index(IndexOpsMixin[S1], PandasObject): +class Index(IndexOpsMixin[S1]): __hash__: ClassVar[None] # type: ignore[assignment] # overloads with additional dtypes @overload diff --git a/pandas-stubs/core/indexes/frozen.pyi b/pandas-stubs/core/indexes/frozen.pyi index e23878fa6..5cebf3ea2 100644 --- a/pandas-stubs/core/indexes/frozen.pyi +++ b/pandas-stubs/core/indexes/frozen.pyi @@ -1,6 +1,4 @@ -from pandas.core.base import PandasObject as PandasObject - -class FrozenList(PandasObject, list): +class FrozenList(list): def union(self, other) -> FrozenList: ... def difference(self, other) -> FrozenList: ... def __getitem__(self, n): ... diff --git a/pandas-stubs/io/sql.pyi b/pandas-stubs/io/sql.pyi index 045982500..a0430e278 100644 --- a/pandas-stubs/io/sql.pyi +++ b/pandas-stubs/io/sql.pyi @@ -11,7 +11,6 @@ from typing import ( overload, ) -from pandas.core.base import PandasObject from pandas.core.frame import DataFrame import sqlalchemy.engine from sqlalchemy.orm import FromStatement @@ -114,7 +113,7 @@ def read_sql( dtype_backend: DtypeBackend | NoDefault = ..., ) -> DataFrame: ... -class PandasSQL(PandasObject): +class PandasSQL: def read_sql(self, *args, **kwargs): ... def to_sql( self, @@ -131,7 +130,7 @@ class PandasSQL(PandasObject): | None = ..., ) -> int | None: ... -class SQLTable(PandasObject): +class SQLTable: name: str pd_sql: PandasSQL # pandas SQL interface prefix: str diff --git a/pandas-stubs/plotting/_core.pyi b/pandas-stubs/plotting/_core.pyi index 303047a2f..768231ff0 100644 --- a/pandas-stubs/plotting/_core.pyi +++ b/pandas-stubs/plotting/_core.pyi @@ -17,7 +17,6 @@ from matplotlib.lines import Line2D import numpy as np import pandas as pd from pandas import Series -from pandas.core.base import PandasObject from pandas.core.frame import DataFrame from scipy.stats.kde import gaussian_kde from typing_extensions import TypeAlias @@ -85,7 +84,7 @@ def boxplot( **kwargs, ) -> _BoxPlotT: ... -class PlotAccessor(PandasObject): +class PlotAccessor: def __init__(self, data) -> None: ... @overload def __call__( From 9ce5fc42bf1b20b76384ee0dce58ea5498df99a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Fri, 4 Aug 2023 04:17:34 -0400 Subject: [PATCH 14/23] works (except for np.ndarray) --- pandas-stubs/core/computation/ops.pyi | 1 - pandas-stubs/core/computation/pytables.pyi | 2 +- pandas-stubs/core/indexes/accessors.pyi | 5 +-- pandas-stubs/core/indexes/base.pyi | 40 ++++++++++++++++++++-- pandas-stubs/core/indexes/timedeltas.pyi | 12 +++++++ pyproject.toml | 2 -- tests/test_timefuncs.py | 2 +- 7 files changed, 53 insertions(+), 11 deletions(-) diff --git a/pandas-stubs/core/computation/ops.pyi b/pandas-stubs/core/computation/ops.pyi index 70ce7351d..3aeecfbd9 100644 --- a/pandas-stubs/core/computation/ops.pyi +++ b/pandas-stubs/core/computation/ops.pyi @@ -34,7 +34,6 @@ class Term: def ndim(self) -> int: ... class Constant(Term): - def __init__(self, value, env, side=..., encoding=...) -> None: ... @property def name(self): ... diff --git a/pandas-stubs/core/computation/pytables.pyi b/pandas-stubs/core/computation/pytables.pyi index 85135c934..e8b5d24c0 100644 --- a/pandas-stubs/core/computation/pytables.pyi +++ b/pandas-stubs/core/computation/pytables.pyi @@ -28,7 +28,7 @@ class Term(ops.Term): def value(self, new_value) -> None: ... class Constant(Term): - def __init__(self, value, env: PyTablesScope, side=..., encoding=...) -> None: ... + def __init__(self, name, env: PyTablesScope, side=..., encoding=...) -> None: ... class BinOp(ops.BinOp): op: str diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index db07d6e99..04890432d 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -343,7 +343,6 @@ class PeriodProperties( _IsLeapYearProperty, _FreqProperty[BaseOffset], ): ... - class CombinedDatetimelikeProperties( DatetimeProperties[ Series[int], @@ -358,9 +357,7 @@ class CombinedDatetimelikeProperties( ], _TimedeltaPropertiesNoRounding[Series[int], Series[float]], _PeriodProperties, -): - def __new__(cls, data: Series): ... - +): ... class TimestampProperties( DatetimeProperties[ Series[int], diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2d7c173bc..65d54d76c 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -64,6 +64,17 @@ class Index(IndexOpsMixin[S1]): __hash__: ClassVar[None] # type: ignore[assignment] # overloads with additional dtypes @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[int | np.integer], + *, + dtype: Literal["int"] | type_t[int] | type_t[np.integer] = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> Index[int]: ... + @overload def __new__( # type: ignore[misc] cls, data: Iterable, @@ -75,6 +86,20 @@ class Index(IndexOpsMixin[S1]): **kwargs, ) -> Index[int]: ... @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[float | np.float32 | np.float64], + *, + dtype: Literal["float"] + | type_t[float] + | type_t[np.float32] + | type_t[np.float64] = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> Index[float]: ... + @overload def __new__( # type: ignore[misc] cls, data: Iterable, @@ -89,6 +114,17 @@ class Index(IndexOpsMixin[S1]): **kwargs, ) -> Index[float]: ... @overload + def __new__( # type: ignore[misc] + cls, + data: Iterable[complex], + *, + dtype: Literal["complex"] | type_t[complex] = ..., + copy: bool = ..., + name=..., + tupleize_cols: bool = ..., + **kwargs, + ) -> Index[complex]: ... + @overload def __new__( # type: ignore[misc] cls, data: Iterable, @@ -194,7 +230,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Iterable[S1] = ..., *, - dtype=..., + dtype: type[S1] = ..., copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -205,7 +241,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Iterable = ..., *, - dtype: type[S1] = ..., + dtype: type[S1], copy: bool = ..., name=..., tupleize_cols: bool = ..., diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 076bf767c..a651b1f83 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -44,6 +44,18 @@ class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties copy: bool = ..., name: str = ..., ) -> None: ... + def __new__( + cls, + data: AnyArrayLike + | list[str] + | Sequence[dt.timedelta | Timedelta | np.timedelta64 | float] = ..., + unit: Literal["D", "h", "m", "s", "ms", "us", "ns"] = ..., + freq: str | BaseOffset = ..., + closed: object = ..., + dtype: Literal[" Self: ... # various ignores needed for mypy, as we do want to restrict what can be used in # arithmetic for these types @overload diff --git a/pyproject.toml b/pyproject.toml index 717e9a13d..d152d61c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -180,7 +180,6 @@ stubPath = "." include = ["tests", "pandas-stubs"] enableTypeIgnoreComments = false # use pyright-specific ignores # disable subset of strict -reportInconsistentConstructor = false reportMissingParameterType = false reportMissingTypeArgument = false reportMissingTypeStubs = false @@ -191,7 +190,6 @@ reportUnknownLambdaType = false reportUnknownMemberType = false reportUnknownParameterType = false reportUnknownVariableType = false -reportUntypedBaseClass = false reportUnusedVariable = false reportPrivateUsage = false # enable optional checks diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 4977d43ca..2b68a0eba 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -768,7 +768,7 @@ def test_to_timedelta_index() -> None: arg3 = tuple(arg1) arg4 = range(0, 10) arg5 = np.arange(10) - arg6 = pd.Index(arg5) + arg6: pd.Index[int] = pd.Index(arg5) check( assert_type(pd.to_timedelta(arg0, "ms"), pd.TimedeltaIndex), pd.TimedeltaIndex ) From d32f53479f91c6cd38d9e0874e74e764fb1ddff0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Tue, 8 Aug 2023 09:49:43 -0400 Subject: [PATCH 15/23] address the easy-to-fix comments --- pandas-stubs/_typing.pyi | 4 ---- pandas-stubs/core/indexes/accessors.pyi | 3 +-- pandas-stubs/core/indexes/multi.pyi | 3 +-- pandas-stubs/core/indexes/timedeltas.pyi | 12 ------------ tests/test_indexes.py | 8 ++++---- 5 files changed, 6 insertions(+), 24 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 4306aede0..7b3f98583 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -520,10 +520,6 @@ S1 = TypeVar( | Interval[Timedelta] | CategoricalDtype, ) -T1 = TypeVar( - "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic] -) -T2 = TypeVar("T2", str, int) IndexingInt: TypeAlias = ( int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8 diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index 04890432d..38284618d 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -36,8 +36,7 @@ from pandas._typing import ( np_ndarray_bool, ) -class Properties(PandasDelegate, NoNewAttributesMixin): - def __init__(self, data: Series, orig) -> None: ... +class Properties(PandasDelegate, NoNewAttributesMixin): ... _DTFieldOpsReturnType = TypeVar("_DTFieldOpsReturnType", Series[int], Index[int]) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index f3afc402a..7f82d0d95 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -15,7 +15,6 @@ from pandas.core.indexes.base import Index from typing_extensions import Self from pandas._typing import ( - T1, Dtype, DtypeArg, HashableT, @@ -156,7 +155,7 @@ class MultiIndex(Index[Any]): def union(self, other, sort=...): ... def intersection(self, other, sort: bool = ...): ... def difference(self, other, sort=...): ... - def astype(self, dtype: DtypeArg | T1, copy: bool = ...) -> Self: ... + def astype(self, dtype: DtypeArg, copy: bool = ...) -> Self: ... def insert(self, loc, item): ... def delete(self, loc): ... def isin(self, values, level=...) -> np_ndarray_bool: ... diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index a651b1f83..bcf480ee3 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -32,18 +32,6 @@ from pandas._typing import ( ) class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties): - def __init__( - self, - data: AnyArrayLike - | list[str] - | Sequence[dt.timedelta | Timedelta | np.timedelta64 | float] = ..., - unit: Literal["D", "h", "m", "s", "ms", "us", "ns"] = ..., - freq: str | BaseOffset = ..., - closed: object = ..., - dtype: Literal[" None: ... def __new__( cls, data: AnyArrayLike diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 627764210..92f385f37 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -69,8 +69,8 @@ def test_multiindex_get_level_values() -> None: def test_index_tolist() -> None: i1 = pd.Index([1, 2, 3]) - check(assert_type(i1.tolist(), "list[int]"), list, int) - check(assert_type(i1.to_list(), "list[int]"), list, int) + check(assert_type(i1.tolist(), list[int]), list, int) + check(assert_type(i1.to_list(), list[int]), list, int) def test_column_getitem() -> None: @@ -708,8 +708,8 @@ def test_interval_index_tuples(): def test_sorted_and_list() -> None: # GH 497 i1 = pd.Index([3, 2, 1]) - check(assert_type(sorted(i1), "list[int]"), list, int) - check(assert_type(list(i1), "list[int]"), list, int) + check(assert_type(sorted(i1), list[int]), list, int) + check(assert_type(list(i1), list[int]), list, int) def test_index_operators() -> None: From 4f1310995eeaa85797f7cfeeba917a591bc404e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Tue, 8 Aug 2023 20:21:59 -0400 Subject: [PATCH 16/23] overloads for numpy --- pandas-stubs/_typing.pyi | 25 ++++++------------- pandas-stubs/core/indexes/base.pyi | 40 +++++++++++++++--------------- pyproject.toml | 1 - tests/test_indexes.py | 12 +++------ tests/test_timefuncs.py | 2 +- 5 files changed, 32 insertions(+), 48 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 7b3f98583..96df2c904 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -511,13 +511,10 @@ S1 = TypeVar( | float | complex | Dtype - | Timestamp - | Timedelta + | datetime.datetime # includes pd.Timestamp + | datetime.timedelta # includes pd.Timedelta | Period - | Interval[int] - | Interval[float] - | Interval[Timestamp] - | Interval[Timedelta] + | Interval[int | float | Timestamp | Timedelta] | CategoricalDtype, ) @@ -597,14 +594,9 @@ ByT = TypeVar( | int | float | complex - | Timestamp - | Timedelta | Scalar | Period - | Interval[int] - | Interval[float] - | Interval[Timestamp] - | Interval[Timedelta] + | Interval[int | float | Timestamp | Timedelta] | tuple, ) # Use a distinct SeriesByT when using groupby with Series of known dtype. @@ -618,13 +610,10 @@ SeriesByT = TypeVar( | int | float | complex - | Timestamp - | Timedelta + | datetime.datetime + | datetime.timedelta | Period - | Interval[int] - | Interval[float] - | Interval[Timestamp] - | Interval[Timedelta], + | Interval[int | float | Timestamp | Timedelta], ) GroupByObjectNonScalar: TypeAlias = ( tuple diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 65d54d76c..8dbc83f73 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -53,6 +53,7 @@ from pandas._typing import ( np_ndarray_anyint, np_ndarray_bool, np_ndarray_int64, + npt, type_t, ) @@ -66,9 +67,9 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[int | np.integer], + data: Sequence[int | np.integer] | IndexOpsMixin[int] | npt.NDArray[np.integer], *, - dtype: Literal["int"] | type_t[int] | type_t[np.integer] = ..., + dtype: Literal["int"] | type_t[int | np.integer] = ..., copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -79,7 +80,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Iterable, *, - dtype: Literal["int"] | type_t[int] | type_t[np.integer], + dtype: Literal["int"] | type_t[int | np.integer], copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -88,12 +89,11 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[float | np.float32 | np.float64], + data: Sequence[float | np.floating] + | IndexOpsMixin[float] + | npt.NDArray[np.floating], *, - dtype: Literal["float"] - | type_t[float] - | type_t[np.float32] - | type_t[np.float64] = ..., + dtype: Literal["float"] | type_t[float | np.floating] = ..., copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -104,10 +104,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Iterable, *, - dtype: Literal["float"] - | type_t[float] - | type_t[np.float32] - | type_t[np.float64], + dtype: Literal["float"] | type_t[float | np.floating], copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -116,9 +113,11 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[complex], + data: Sequence[complex | np.complexfloating] + | IndexOpsMixin[complex] + | npt.NDArray[np.complexfloating], *, - dtype: Literal["complex"] | type_t[complex] = ..., + dtype: Literal["complex"] | type_t[complex | np.complexfloating] = ..., copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -129,7 +128,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Iterable, *, - dtype: Literal["complex"] | type_t[complex], + dtype: Literal["complex"] | type_t[complex | np.complexfloating], copy: bool = ..., name=..., tupleize_cols: bool = ..., @@ -139,7 +138,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[np.datetime64 | datetime], + data: Sequence[np.datetime64 | datetime] | IndexOpsMixin[datetime], *, dtype: TimestampDtypeArg = ..., copy: bool = ..., @@ -161,7 +160,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[Period], + data: Sequence[Period] | IndexOpsMixin[Period], *, dtype: PeriodDtype = ..., copy: bool = ..., @@ -183,7 +182,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[np.timedelta64 | timedelta], + data: Sequence[np.timedelta64 | timedelta] | IndexOpsMixin[timedelta], *, dtype: TimedeltaDtypeArg = ..., copy: bool = ..., @@ -205,7 +204,8 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Iterable[Interval[_OrderableT]], + data: Sequence[Interval[_OrderableT]] # type: ignore[type-var] + | IndexOpsMixin[Interval[_OrderableT]], *, dtype: Literal["Interval"] = ..., copy: bool = ..., @@ -228,7 +228,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( cls, - data: Iterable[S1] = ..., + data: Iterable[S1] | IndexOpsMixin[S1] = ..., *, dtype: type[S1] = ..., copy: bool = ..., diff --git a/pyproject.toml b/pyproject.toml index d152d61c5..c1970821e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -195,7 +195,6 @@ reportPrivateUsage = false # enable optional checks reportMissingModuleSource = true useLibraryCodeForTypes = false -defineConstant = { MYPY_CHECKING = false } [tool.codespell] ignore-words-list = "indext, mose, sav, ser" diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 92f385f37..801e79c5d 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -2,7 +2,6 @@ import datetime as dt from typing import ( - TYPE_CHECKING, Tuple, Union, ) @@ -22,9 +21,6 @@ check, ) -if TYPE_CHECKING: - MYPY_CHECKING: bool = True - def test_index_unique() -> None: df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2])) @@ -69,8 +65,8 @@ def test_multiindex_get_level_values() -> None: def test_index_tolist() -> None: i1 = pd.Index([1, 2, 3]) - check(assert_type(i1.tolist(), list[int]), list, int) - check(assert_type(i1.to_list(), list[int]), list, int) + check(assert_type(i1.tolist(), "list[int]"), list, int) + check(assert_type(i1.to_list(), "list[int]"), list, int) def test_column_getitem() -> None: @@ -708,8 +704,8 @@ def test_interval_index_tuples(): def test_sorted_and_list() -> None: # GH 497 i1 = pd.Index([3, 2, 1]) - check(assert_type(sorted(i1), list[int]), list, int) - check(assert_type(list(i1), list[int]), list, int) + check(assert_type(sorted(i1), "list[int]"), list, int) + check(assert_type(list(i1), "list[int]"), list, int) def test_index_operators() -> None: diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 2b68a0eba..4977d43ca 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -768,7 +768,7 @@ def test_to_timedelta_index() -> None: arg3 = tuple(arg1) arg4 = range(0, 10) arg5 = np.arange(10) - arg6: pd.Index[int] = pd.Index(arg5) + arg6 = pd.Index(arg5) check( assert_type(pd.to_timedelta(arg0, "ms"), pd.TimedeltaIndex), pd.TimedeltaIndex ) From bd3962532439f418b53b81c51188a02e46d61a5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Tue, 8 Aug 2023 20:36:44 -0400 Subject: [PATCH 17/23] did numexpr break the CI? --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index c1970821e..2ab7d6185 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ jinja2 = ">=3.1" scipy = ">=1.9.1" SQLAlchemy = ">=2.0.12" types-python-dateutil = ">=2.8.19" +numexpr = "<2.8.5" [build-system] requires = ["poetry-core>=1.0.0"] From 8c48567240daccd84998e91a47eaf30473fd83c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Wed, 9 Aug 2023 10:01:59 -0400 Subject: [PATCH 18/23] lie: df.columns -> Index[str] --- pandas-stubs/core/frame.pyi | 2 +- tests/test_frame.py | 5 ++--- tests/test_indexes.py | 8 ++++---- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index a8f46facb..caa02427a 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1482,7 +1482,7 @@ class DataFrame(NDFrame, OpsMixin): @property def at(self): ... # Not sure what to do with this yet; look at source @property - def columns(self) -> Index: ... + def columns(self) -> Index[str]: ... @columns.setter # setter needs to be right next to getter; otherwise mypy complains def columns( self, cols: AnyArrayLike | list[HashableT] | tuple[HashableT, ...] diff --git a/tests/test_frame.py b/tests/test_frame.py index 46a689aa2..e577be021 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1761,8 +1761,7 @@ def test_getmultiindex_columns() -> None: [(i, s) for i in [1] for s in df.columns.get_level_values(1)] ] res4: pd.DataFrame = df[[df.columns[0]]] - column: Scalar = df.columns[0] - check(assert_type(df[column], pd.Series), pd.Series) + check(assert_type(df[df.columns[0]], pd.Series), pd.Series) check(assert_type(df[li[0]], pd.Series), pd.Series) @@ -2592,7 +2591,7 @@ def test_in_columns() -> None: # GH 532 (PR) df = pd.DataFrame(np.random.random((3, 4)), columns=["cat", "dog", "rat", "pig"]) cols = [c for c in df.columns if "at" in c] - check(assert_type(cols, list), list, str) + check(assert_type(cols, "list[str]"), list, str) check(assert_type(df.loc[:, cols], pd.DataFrame), pd.DataFrame) check(assert_type(df[cols], pd.DataFrame), pd.DataFrame) check(assert_type(df.groupby(by=cols).sum(), pd.DataFrame), pd.DataFrame) diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 801e79c5d..be437f5e0 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -73,9 +73,9 @@ def test_column_getitem() -> None: # https://github.com/microsoft/python-type-stubs/issues/199#issuecomment-1132806594 df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"]) - columns: pd.Index[str] = df.columns - check(assert_type(columns[0], str), str) - check(assert_type(df[columns[0]], pd.Series), pd.Series, np.int64) + column = df.columns[0] + check(assert_type(column, str), str) + check(assert_type(df[column], pd.Series), pd.Series, np.int64) def test_column_contains() -> None: @@ -93,7 +93,7 @@ def test_column_sequence() -> None: df = pd.DataFrame([1, 2, 3]) col_list = list(df.columns) check( - assert_type(col_list, list), + assert_type(col_list, "list[str]"), list, int, ) From 67f5e8de911ac6a6db627154158e277eb7d4c446 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Wed, 9 Aug 2023 10:06:48 -0400 Subject: [PATCH 19/23] new ruff has far more pyi rules --- .pre-commit-config.yaml | 2 +- pandas-stubs/_typing.pyi | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1786d3276..7aabe629a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,7 @@ repos: hooks: - id: isort - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.282 + rev: v0.0.283 hooks: - id: ruff args: [ diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 96df2c904..7febda0c5 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -91,7 +91,7 @@ DtypeBackend: TypeAlias = Literal["pyarrow", "numpy_nullable"] BooleanDtypeArg: TypeAlias = ( # Builtin bool type and its string alias - type[bool] # noqa: PYI030 + type[bool] # noqa: PYI030,PYI055 | Literal["bool"] # Pandas nullable boolean type and its string alias | pd.BooleanDtype @@ -105,7 +105,7 @@ BooleanDtypeArg: TypeAlias = ( ) IntDtypeArg: TypeAlias = ( # Builtin integer type and its string alias - type[int] # noqa: PYI030 + type[int] # noqa: PYI030,PYI055 | Literal["int"] # Pandas nullable integer types and their string aliases | pd.Int8Dtype @@ -137,7 +137,7 @@ IntDtypeArg: TypeAlias = ( ) UIntDtypeArg: TypeAlias = ( # Pandas nullable unsigned integer types and their string aliases - pd.UInt8Dtype # noqa: PYI030 + pd.UInt8Dtype # noqa: PYI030,PYI055 | pd.UInt16Dtype | pd.UInt32Dtype | pd.UInt64Dtype @@ -166,7 +166,7 @@ UIntDtypeArg: TypeAlias = ( ) FloatDtypeArg: TypeAlias = ( # Builtin float type and its string alias - type[float] # noqa: PYI030 + type[float] # noqa: PYI030,PYI055 | Literal["float"] # Pandas nullable float types and their string aliases | pd.Float32Dtype @@ -197,7 +197,7 @@ FloatDtypeArg: TypeAlias = ( ) ComplexDtypeArg: TypeAlias = ( # Builtin complex type and its string alias - type[complex] # noqa: PYI030 + type[complex] # noqa: PYI030,PYI055 | Literal["complex"] # Numpy complex types and their aliases # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.csingle @@ -326,7 +326,7 @@ TimestampDtypeArg: TypeAlias = Literal[ StrDtypeArg: TypeAlias = ( # Builtin str type and its string alias - type[str] # noqa: PYI030 + type[str] # noqa: PYI030,PYI055 | Literal["str"] # Pandas nullable string type and its string alias | pd.StringDtype @@ -340,7 +340,7 @@ StrDtypeArg: TypeAlias = ( ) BytesDtypeArg: TypeAlias = ( # Builtin bytes type and its string alias - type[bytes] # noqa: PYI030 + type[bytes] # noqa: PYI030,PYI055 | Literal["bytes"] # Numpy bytes type and its string alias # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bytes_ @@ -353,7 +353,7 @@ CategoryDtypeArg: TypeAlias = CategoricalDtype | Literal["category"] ObjectDtypeArg: TypeAlias = ( # Builtin object type and its string alias - type[object] # noqa: PYI030 + type[object] # noqa: PYI030,PYI055 | Literal["object"] # Numpy object type and its string alias # https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.object_ From 566c91882ce56c4fb385ae6716825f04ce7bcb02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 10 Aug 2023 07:58:18 -0400 Subject: [PATCH 20/23] Make it clear that both mypy&pyright infer Never; but only pyright detects it as invalid --- pyproject.toml | 2 +- tests/test_scalars.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2ab7d6185..b7cba1eda 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,7 @@ jinja2 = ">=3.1" scipy = ">=1.9.1" SQLAlchemy = ">=2.0.12" types-python-dateutil = ">=2.8.19" -numexpr = "<2.8.5" +numexpr = "<2.8.5" # https://github.com/pandas-dev/pandas/issues/54449 [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 1ef0e543f..ad0bcd2ec 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -747,10 +747,12 @@ def test_timedelta_mul_div() -> None: mp_series_int // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] assert_type( - md_int64_index // td, Never # pyright: ignore[reportGeneralTypeIssues] + md_int64_index // td, # pyright: ignore[reportGeneralTypeIssues] + Never, ) assert_type( - md_float_index // td, Never # pyright: ignore[reportGeneralTypeIssues] + md_float_index // td, # pyright: ignore[reportGeneralTypeIssues] + Never, ) check(assert_type(td / td, float), float) @@ -785,10 +787,12 @@ def test_timedelta_mul_div() -> None: mp_series_int / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] assert_type( - md_int64_index / td, Never # pyright: ignore[reportGeneralTypeIssues] + md_int64_index / td, # pyright: ignore[reportGeneralTypeIssues] + Never, ) assert_type( - md_float_index / td, Never # pyright: ignore[reportGeneralTypeIssues] + md_float_index / td, # pyright: ignore[reportGeneralTypeIssues] + Never, ) From ff5ee587795bb7c92af6ea8472f474b6c3a8dd42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Thu, 10 Aug 2023 10:32:52 -0400 Subject: [PATCH 21/23] use type aliases instead of npt --- pandas-stubs/_typing.pyi | 2 ++ pandas-stubs/core/indexes/base.pyi | 11 +++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 7febda0c5..3ee9134d0 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -483,6 +483,8 @@ ScalarT = TypeVar("ScalarT", bound=Scalar) np_ndarray_int64: TypeAlias = npt.NDArray[np.int64] np_ndarray_int: TypeAlias = npt.NDArray[np.signedinteger] np_ndarray_anyint: TypeAlias = npt.NDArray[np.integer] +np_ndarray_float: TypeAlias = npt.NDArray[np.floating] +np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating] np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_] np_ndarray_str: TypeAlias = npt.NDArray[np.str_] diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 8dbc83f73..f12d7ff22 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -52,8 +52,9 @@ from pandas._typing import ( TimestampDtypeArg, np_ndarray_anyint, np_ndarray_bool, + np_ndarray_complex, + np_ndarray_float, np_ndarray_int64, - npt, type_t, ) @@ -67,7 +68,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Sequence[int | np.integer] | IndexOpsMixin[int] | npt.NDArray[np.integer], + data: Sequence[int | np.integer] | IndexOpsMixin[int] | np_ndarray_anyint, *, dtype: Literal["int"] | type_t[int | np.integer] = ..., copy: bool = ..., @@ -89,9 +90,7 @@ class Index(IndexOpsMixin[S1]): @overload def __new__( # type: ignore[misc] cls, - data: Sequence[float | np.floating] - | IndexOpsMixin[float] - | npt.NDArray[np.floating], + data: Sequence[float | np.floating] | IndexOpsMixin[float] | np_ndarray_float, *, dtype: Literal["float"] | type_t[float | np.floating] = ..., copy: bool = ..., @@ -115,7 +114,7 @@ class Index(IndexOpsMixin[S1]): cls, data: Sequence[complex | np.complexfloating] | IndexOpsMixin[complex] - | npt.NDArray[np.complexfloating], + | np_ndarray_complex, *, dtype: Literal["complex"] | type_t[complex | np.complexfloating] = ..., copy: bool = ..., From c8d201972afa8e1a9c5129a886f207cee4a9cc05 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sun, 13 Aug 2023 00:09:05 -0400 Subject: [PATCH 22/23] fix issue with floordiv for mypy --- pandas-stubs/core/arraylike.pyi | 4 +++- pandas-stubs/core/frame.pyi | 4 ++++ pandas-stubs/core/indexes/base.pyi | 14 ++++++++++---- pandas-stubs/core/series.pyi | 2 +- tests/test_scalars.py | 10 ++-------- 5 files changed, 20 insertions(+), 14 deletions(-) diff --git a/pandas-stubs/core/arraylike.pyi b/pandas-stubs/core/arraylike.pyi index c1148a1b4..967b51ed3 100644 --- a/pandas-stubs/core/arraylike.pyi +++ b/pandas-stubs/core/arraylike.pyi @@ -31,7 +31,9 @@ class OpsMixin: def __rmul__(self, other: Any) -> Self: ... def __truediv__(self, other: Any) -> Self: ... def __rtruediv__(self, other: Any) -> Self: ... - def __floordiv__(self, other: Any) -> Self: ... + # __floordiv__ is handled by superclasses that specify only the valid values + # that can be passed + # def __floordiv__(self, other: Any) -> Self: ... def __rfloordiv__(self, other: Any) -> Self: ... def __mod__(self, other: Any) -> Self: ... def __rmod__(self, other: Any) -> Self: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index caa02427a..71e00c245 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -2275,3 +2275,7 @@ class DataFrame(NDFrame, OpsMixin): level: Level | None = ..., drop_level: _bool = ..., ) -> DataFrame | Series: ... + # floordiv overload + def __floordiv__( + self, other: float | DataFrame | Series[int] | Series[float] + ) -> Self: ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index f12d7ff22..95f1277e7 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -427,10 +427,16 @@ class Index(IndexOpsMixin[S1]): ) -> TimedeltaIndex: ... @overload def __mul__(self, other: Any) -> Self: ... - @overload - def __floordiv__(self: Index[int] | Index[float], other: timedelta) -> Never: ... - @overload - def __floordiv__(self, other: Any) -> Self: ... + def __floordiv__( + self, + other: float + | Series[int] + | Series[float] + | Sequence[int] + | Sequence[float] + | Index[int] + | Index[float], + ) -> Self: ... @overload def __truediv__(self: Index[int] | Index[float], other: timedelta) -> Never: ... @overload diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index df9db7f42..839c2a19a 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1450,7 +1450,7 @@ class Series(IndexOpsMixin[S1], NDFrame): # def __array__(self, dtype: Optional[_bool] = ...) -> _np_ndarray def __div__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __eq__(self, other: object) -> Series[_bool]: ... # type: ignore[override] - def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ... # type: ignore[override] + def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ... def __ge__( # type: ignore[override] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta ) -> Series[_bool]: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index ad0bcd2ec..652814fec 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -746,14 +746,8 @@ def test_timedelta_mul_div() -> None: md_ndarray_float // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] mp_series_int // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - assert_type( - md_int64_index // td, # pyright: ignore[reportGeneralTypeIssues] - Never, - ) - assert_type( - md_float_index // td, # pyright: ignore[reportGeneralTypeIssues] - Never, - ) + md_int64_index // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + md_float_index // td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] check(assert_type(td / td, float), float) check(assert_type(td / pd.NaT, float), float) From 585c93b79d670db6133513dae8da20a9ddc7d700 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sun, 13 Aug 2023 00:11:09 -0400 Subject: [PATCH 23/23] fix comment in arraylike --- pandas-stubs/core/arraylike.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas-stubs/core/arraylike.pyi b/pandas-stubs/core/arraylike.pyi index 967b51ed3..9c8d95f43 100644 --- a/pandas-stubs/core/arraylike.pyi +++ b/pandas-stubs/core/arraylike.pyi @@ -31,7 +31,7 @@ class OpsMixin: def __rmul__(self, other: Any) -> Self: ... def __truediv__(self, other: Any) -> Self: ... def __rtruediv__(self, other: Any) -> Self: ... - # __floordiv__ is handled by superclasses that specify only the valid values + # __floordiv__ is handled by subclasses that specify only the valid values # that can be passed # def __floordiv__(self, other: Any) -> Self: ... def __rfloordiv__(self, other: Any) -> Self: ...