From 12505bec11cc7c756a9048a41ca7a56c59d00338 Mon Sep 17 00:00:00 2001 From: Amotz Date: Fri, 26 Aug 2022 20:16:34 +0300 Subject: [PATCH 1/3] Cleaned uptyping on __new__ for Series and DataFrame --- pandas-stubs/core/frame.pyi | 14 ++++++++++++-- pandas-stubs/core/series.pyi | 12 ++++++------ 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 0c397e46f..2bd822d48 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -170,6 +170,7 @@ class DataFrame(NDFrame, OpsMixin): __hash__: ClassVar[None] # type: ignore[assignment] + @overload def __new__( cls, data: ListLikeU @@ -177,8 +178,17 @@ class DataFrame(NDFrame, OpsMixin): | dict[Any, Any] | Iterable[tuple[Hashable, ListLikeU]] | None = ..., - index: Axes | None = ..., - columns: Axes | None = ..., + index: Iterable[Hashable] | None = ..., + columns: Iterable[Hashable] | None = ..., + dtype=..., + copy: _bool = ..., + ) -> DataFrame: ... + @overload + def __new__( + cls, + data: Scalar, + index: Iterable[Hashable], + columns: Iterable[Hashable], dtype=..., copy: _bool = ..., ) -> DataFrame: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 07018808d..59280727b 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -122,14 +122,14 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): idx: MaskType | Index | Sequence[float] - | list[_str] + | list[str] | slice | tuple[str | float | slice | Index, ...], ) -> Series[S1]: ... @overload def __getitem__( self, - idx: _str | float, + idx: str | float, ) -> S1: ... @overload def __setitem__( @@ -159,7 +159,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __new__( cls, data: DatetimeIndex, - index: _str | int | Series | list | Index | None = ..., + index: Iterable[Hashable] | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., @@ -169,7 +169,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __new__( cls, data: PeriodIndex, - index: _str | int | Series | list | Index | None = ..., + index: Iterable[Hashable] | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., @@ -180,7 +180,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): cls, data: object | _ListLike | Series[S1] | dict[int, S1] | dict[_str, S1] | None, dtype: type[S1], - index: _str | int | Series | list | Index | None = ..., + index: Iterable[Hashable] | None = ..., name: Hashable | None = ..., copy: bool = ..., fastpath: bool = ..., @@ -194,7 +194,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): | dict[int, S1] | dict[_str, S1] | None = ..., - index: _str | int | Series | list | Index | None = ..., + index: Iterable[Hashable] | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., From 21aad859edb146fc35c252a20ef640f7a30c0973 Mon Sep 17 00:00:00 2001 From: Amotz Date: Fri, 26 Aug 2022 21:41:52 +0300 Subject: [PATCH 2/3] Added tests --- pandas-stubs/core/frame.pyi | 8 ++++---- pandas-stubs/core/series.pyi | 8 ++++---- tests/test_frame.py | 1 + 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 2bd822d48..ca8d7bbdd 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -178,8 +178,8 @@ class DataFrame(NDFrame, OpsMixin): | dict[Any, Any] | Iterable[tuple[Hashable, ListLikeU]] | None = ..., - index: Iterable[Hashable] | None = ..., - columns: Iterable[Hashable] | None = ..., + index: Axes | None = ..., + columns: Axes | None = ..., dtype=..., copy: _bool = ..., ) -> DataFrame: ... @@ -187,8 +187,8 @@ class DataFrame(NDFrame, OpsMixin): def __new__( cls, data: Scalar, - index: Iterable[Hashable], - columns: Iterable[Hashable], + index: Axes, + columns: Axes, dtype=..., copy: _bool = ..., ) -> DataFrame: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 59280727b..cf45b3e5d 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -159,7 +159,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __new__( cls, data: DatetimeIndex, - index: Iterable[Hashable] | None = ..., + index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., @@ -169,7 +169,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __new__( cls, data: PeriodIndex, - index: Iterable[Hashable] | None = ..., + index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., @@ -180,7 +180,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): cls, data: object | _ListLike | Series[S1] | dict[int, S1] | dict[_str, S1] | None, dtype: type[S1], - index: Iterable[Hashable] | None = ..., + index: Axes | None = ..., name: Hashable | None = ..., copy: bool = ..., fastpath: bool = ..., @@ -194,7 +194,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): | dict[int, S1] | dict[_str, S1] | None = ..., - index: Iterable[Hashable] | None = ..., + index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index 8aed66021..bc3f42d8d 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -42,6 +42,7 @@ def test_types_init() -> None: dtype=np.int8, copy=True, ) + pd.DataFrame(0, index=[0, 1], columns=[0, 1]) def test_types_all() -> None: From 6a358f17b4d9b12c2aee2bdfb3bae2df398e73ec Mon Sep 17 00:00:00 2001 From: Amotz Date: Fri, 26 Aug 2022 22:53:29 +0300 Subject: [PATCH 3/3] Added checkoing to the test --- tests/test_frame.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index bc3f42d8d..1a37e083a 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -42,7 +42,10 @@ def test_types_init() -> None: dtype=np.int8, copy=True, ) - pd.DataFrame(0, index=[0, 1], columns=[0, 1]) + check( + assert_type(pd.DataFrame(0, index=[0, 1], columns=[0, 1]), pd.DataFrame), + pd.DataFrame, + ) def test_types_all() -> None: