Seris and DataFrame creation typing (#230)

* Cleaned uptyping on __new__ for Series and DataFrame

* Added tests

* Added checkoing to the test

Co-authored-by: Amotz <[email protected]>

Author: amotzop

pandas-stubs/core/frame.pyi

@@ -171,6 +171,7 @@ class DataFrame(NDFrame, OpsMixin):
 
     __hash__: ClassVar[None]  # type: ignore[assignment]
 
+    @overload
     def __new__(
         cls,
         data: ListLikeU
@@ -183,6 +184,15 @@ class DataFrame(NDFrame, OpsMixin):
         dtype=...,
         copy: _bool = ...,
     ) -> DataFrame: ...
+    @overload
+    def __new__(
+        cls,
+        data: Scalar,
+        index: Axes,
+        columns: Axes,
+        dtype=...,
+        copy: _bool = ...,
+    ) -> DataFrame: ...
     @property
     def axes(self) -> list[Index]: ...
     @property

pandas-stubs/core/series.pyi

@@ -124,14 +124,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__(
@@ -161,7 +161,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     def __new__(
         cls,
         data: DatetimeIndex,
-        index: _str | int | Series | list | Index | None = ...,
+        index: Axes | None = ...,
         dtype=...,
         name: Hashable | None = ...,
         copy: bool = ...,
@@ -171,7 +171,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     def __new__(
         cls,
         data: PeriodIndex,
-        index: _str | int | Series | list | Index | None = ...,
+        index: Axes | None = ...,
         dtype=...,
         name: Hashable | None = ...,
         copy: bool = ...,
@@ -182,7 +182,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: Axes | None = ...,
         name: Hashable | None = ...,
         copy: bool = ...,
         fastpath: bool = ...,
@@ -196,7 +196,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         | dict[int, S1]
         | dict[_str, S1]
         | None = ...,
-        index: _str | int | Series | list | Index | None = ...,
+        index: Axes | None = ...,
         dtype=...,
         name: Hashable | None = ...,
         copy: bool = ...,

tests/test_frame.py

@@ -42,6 +42,10 @@ def test_types_init() -> None:
         dtype=np.int8,
         copy=True,
     )
+    check(
+        assert_type(pd.DataFrame(0, index=[0, 1], columns=[0, 1]), pd.DataFrame),
+        pd.DataFrame,
+    )
 
 
 def test_types_all() -> None: