Allow subtypes of List[Scalar] by introducing ScalarArg (#124)

* Allow subtypes of List[Scalar] by introducing ScalarArg

* add check on result

* black on test_frame.py

* ScalarArg -> ScalarT

Author: Dr-Irv

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]

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,

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)