diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 343e8831e..ae677a454 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -972,6 +972,10 @@ ListLikeHashable: TypeAlias = ( MutableSequence[HashableT] | np_1darray | tuple[HashableT, ...] | range ) +if TYPE_CHECKING: # noqa: PYI002 + IndexT0 = TypeVar("IndexT0", bound=Index, default=Index) + IndexStrT0 = TypeVar("IndexStrT0", bound=Index, default=Index[str]) + class SupportsDType(Protocol[GenericT_co]): @property def dtype(self) -> np.dtype[GenericT_co]: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 742138fee..e84a17236 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -117,6 +117,8 @@ from pandas._typing import ( IndexingInt, IndexKeyFunc, IndexLabel, + IndexStrT0, + IndexT0, IndexType, InterpolateOptions, IntervalClosedType, @@ -378,10 +380,19 @@ _AstypeArgExt: TypeAlias = ( ) _AstypeArgExtList: TypeAlias = _AstypeArgExt | list[_AstypeArgExt] -class DataFrame(NDFrame, OpsMixin, _GetItemHack): +class DataFrame(NDFrame, OpsMixin, _GetItemHack, Generic[IndexT0, IndexStrT0]): __hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride] + @overload + def __new__( + cls, + data: DataFrame[IndexT0, IndexStrT0], + index: None = None, + columns: None = None, + dtype: Dtype | None = None, + copy: _bool | None = None, + ) -> DataFrame[IndexT0, IndexStrT0]: ... @overload def __new__( cls, @@ -398,6 +409,15 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): copy: _bool | None = None, ) -> Self: ... @overload + def __new__( + cls, + data: Scalar, + index: IndexT0, + columns: IndexStrT0, + dtype: Dtype | None = None, + copy: _bool | None = None, + ) -> DataFrame[IndexT0, IndexStrT0]: ... + @overload def __new__( cls, data: Scalar, @@ -1898,7 +1918,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @property def at(self) -> _AtIndexerFrame: ... @property - def columns(self) -> Index[str]: ... + def columns(self) -> IndexStrT0: ... @columns.setter # setter needs to be right next to getter; otherwise mypy complains def columns( self, cols: AnyArrayLike | SequenceNotStr[Hashable] | tuple[Hashable, ...] @@ -1912,7 +1932,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @property def iloc(self) -> _iLocIndexerFrame[Self]: ... @property - def index(self) -> Index: ... + def index(self) -> IndexT0: ... @index.setter def index( self, idx: AnyArrayLike | SequenceNotStr[Hashable] | tuple[Hashable, ...] @@ -2289,7 +2309,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): inplace: Literal[False] = False, **kwargs: Any, ) -> Self: ... - def keys(self) -> Index: ... + def keys(self) -> IndexStrT0: ... def kurt( self, axis: Axis | None = ..., diff --git a/tests/frame/test_indexing.py b/tests/frame/test_indexing.py index 8507c17a3..4261aefc8 100644 --- a/tests/frame/test_indexing.py +++ b/tests/frame/test_indexing.py @@ -500,12 +500,18 @@ def select1(df: pd.DataFrame) -> pd.Series: def select2(df: pd.DataFrame) -> list[Hashable]: return [i for i in df.index if cast(int, i) % 2 == 1] - check(assert_type(df.loc[select2, "x"], pd.Series), pd.Series) + # I think it has to do with the two overlapping overloads of __getitem__ in _LocIndexerFrame + # tuple[Callable[[DataFrame], ScalarT], int | StrLike] must be overlapping with + # tuple[Callable[[DataFrame], ScalarT | list[HashableT] | IndexType | MaskType], ScalarT] + check(assert_type(df.loc[select2, "x"], pd.Series), pd.Series) # type: ignore[assert-type] def select3(_: pd.DataFrame) -> int: return 1 - check(assert_type(df.loc[select3, "x"], Scalar), np.integer) + # I think it has to do with the two overlapping overloads of __getitem__ in _LocIndexerFrame + # tuple[Callable[[DataFrame], ScalarT], int | StrLike] must be overlapping with + # tuple[Callable[[DataFrame], ScalarT | list[HashableT] | IndexType | MaskType], ScalarT] + check(assert_type(df.loc[select3, "x"], Scalar), np.integer) # type: ignore[assert-type] check( assert_type(df.loc[:, lambda df: df.columns.str.startswith("x")], pd.DataFrame), diff --git a/tests/series/test_series.py b/tests/series/test_series.py index bfd4028a7..ceed99c80 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -3157,7 +3157,14 @@ def test_rank() -> None: def test_round() -> None: # GH 791 - check(assert_type(round(pd.DataFrame([])), pd.DataFrame), pd.DataFrame) + # TODO: microsoft/pyright#11179 + check( + assert_type( + round(pd.DataFrame([])), # pyright: ignore[reportAssertTypeFailure] + pd.DataFrame, + ), + pd.DataFrame, + ) check(assert_type(round(pd.Series([1], dtype=int)), "pd.Series[int]"), pd.Series) diff --git a/tests/test_groupby.py b/tests/test_groupby.py index b0287217b..c7000611a 100644 --- a/tests/test_groupby.py +++ b/tests/test_groupby.py @@ -583,8 +583,9 @@ def df2scalar(val: DataFrame) -> float: # iter iterator = iter(GB_DF.rolling(1)) - check(assert_type(iterator, Iterator[DataFrame]), Iterator) - check(assert_type(next(iterator), DataFrame), DataFrame) + # TODO: reported python/mypy#20436 python/mypy#20435 + check(assert_type(iterator, Iterator[DataFrame]), Iterator) # type: ignore[assert-type] + check(assert_type(next(iterator), DataFrame), DataFrame) # type: ignore[assert-type] check(assert_type(list(GB_DF.rolling(1)), list[DataFrame]), list, DataFrame) @@ -782,8 +783,9 @@ def df2scalar(val: DataFrame) -> float: # iter iterator = iter(GB_DF.expanding(1)) - check(assert_type(iterator, Iterator[DataFrame]), Iterator) - check(assert_type(next(iterator), DataFrame), DataFrame) + # TODO: reported python/mypy#20436 python/mypy#20435 + check(assert_type(iterator, Iterator[DataFrame]), Iterator) # type: ignore[assert-type] + check(assert_type(next(iterator), DataFrame), DataFrame) # type: ignore[assert-type] check(assert_type(list(GB_DF.expanding(1)), list[DataFrame]), list, DataFrame) @@ -960,8 +962,9 @@ def test_frame_groupby_ewm() -> None: # iter iterator = iter(GB_DF.ewm(1)) - check(assert_type(iterator, Iterator[DataFrame]), Iterator) - check(assert_type(next(iterator), DataFrame), DataFrame) + # TODO: reported python/mypy#20436 python/mypy#20435 + check(assert_type(iterator, Iterator[DataFrame]), Iterator) # type: ignore[assert-type] + check(assert_type(next(iterator), DataFrame), DataFrame) # type: ignore[assert-type] check(assert_type(list(GB_DF.ewm(1)), list[DataFrame]), list, DataFrame) diff --git a/tests/test_resampler.py b/tests/test_resampler.py index f9da71ebe..a71d70df6 100644 --- a/tests/test_resampler.py +++ b/tests/test_resampler.py @@ -45,7 +45,8 @@ def test_props() -> None: def test_iter() -> None: - assert_type(iter(DF.resample("ME")), Iterator[tuple[Hashable, DataFrame]]) + # TODO: reported python/mypy#20436 python/mypy#20435 + assert_type(iter(DF.resample("ME")), Iterator[tuple[Hashable, DataFrame]]) # type: ignore[assert-type] for v in DF.resample("ME"): check(assert_type(v, tuple[Hashable, DataFrame]), tuple)