diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 8eafcda24..0fc1d674f 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -122,6 +122,7 @@ Scalar = Union[ Timestamp, Timedelta, ] +ScalarT = TypeVar("ScalarT", bound=Scalar) # Refine the definitions below in 3.9 to use the specialized type. np_ndarray_int8 = npt.NDArray[np.int8] np_ndarray_int16 = npt.NDArray[np.int16] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 39f3f2384..610f3f420 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -65,6 +65,7 @@ from pandas._typing import ( MaskType, Renamer, Scalar, + ScalarT, SeriesAxisType, StrLike, T as TType, @@ -368,7 +369,7 @@ class DataFrame(NDFrame, OpsMixin): Series[_bool], DataFrame, List[_str], - List[Scalar], + List[ScalarT], Index, np_ndarray_str, np_ndarray_bool, diff --git a/tests/test_frame.py b/tests/test_frame.py index a50950778..38760e986 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1169,3 +1169,11 @@ def test_func(h: Hashable): test_func(pd.DataFrame()) # type: ignore[arg-type] test_func(pd.Series([], dtype=object)) # type: ignore[arg-type] test_func(pd.Index([])) # type: ignore[arg-type] + + +def test_columns_mixlist() -> None: + # GH 97 + df = pd.DataFrame({"a": [1, 2, 3], 1: [3, 4, 5]}) + key: List[Union[int, str]] + key = [1] + check(assert_type(df[key], pd.DataFrame), pd.DataFrame)