use Union for ListLike in DataFrame __new__ (#198)

Author: Dr-Irv

pandas-stubs/_typing.pyi

@@ -99,6 +99,7 @@ AxisType = Literal["columns", "index", 0, 1]
 DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
 KeysArgType = Any
 ListLike = TypeVar("ListLike", Sequence, np.ndarray, "Series", "Index")
+ListLikeU = Union[Sequence, np.ndarray, Series, Index]
 StrLike = Union[str, np.str_]
 Scalar = Union[
     str,

pandas-stubs/core/frame.pyi

@@ -61,6 +61,7 @@ from pandas._typing import (
     Label,
     Level,
     ListLike,
+    ListLikeU,
     MaskType,
     Renamer,
     Scalar,
@@ -157,10 +158,10 @@ class DataFrame(NDFrame, OpsMixin):
 
     def __new__(
         cls,
-        data: ListLike
+        data: ListLikeU
         | DataFrame
         | dict[Any, Any]
-        | Iterable[tuple[Hashable, ListLike]]
+        | Iterable[tuple[Hashable, ListLikeU]]
         | None = ...,
         index: Axes | None = ...,
         columns: Axes | None = ...,

tests/test_frame.py

@@ -8,10 +8,12 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Generic,
     Hashable,
     Iterable,
     Iterator,
     Tuple,
+    TypeVar,
     Union,
 )
 
@@ -1614,3 +1616,16 @@ def test_dict_items() -> None:
     # GH 180
     x = {"a": [1]}
     check(assert_type(pd.DataFrame(x.items()), pd.DataFrame), pd.DataFrame)
+
+
+def test_generic() -> None:
+    # GH 197
+    T = TypeVar("T")
+
+    class MyDataFrame(pd.DataFrame, Generic[T]):
+        ...
+
+    def func() -> MyDataFrame[int]:
+        return MyDataFrame[int]({"foo": [1, 2, 3]})
+
+    func()