Read-only (covariant) list parameter annotations (#1745)

* Add failing test

* Add covariant list for use in parameter annotations

* Fix hashability and pyright comment moved by black

* type-check only

* Address code review

Author: hamdanal

pandas-stubs/_typing.pyi

@@ -17,6 +17,7 @@ import sys
 from typing import (
     TYPE_CHECKING,
     Any,
+    ClassVar,
     Generic,
     Literal,
     Protocol,
@@ -1230,6 +1231,21 @@ class Just(Protocol, Generic[T]):
     @override
     def __class__(self, t: type[T], /) -> None: ...
 
+# Read-only (covariant) list for use in parameter annotations (See GH #1745)
+class CovariantList(Protocol[_T_co]):
+    __hash__: ClassVar[None]  # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
+    @property  # type: ignore[override]
+    def __class__(self) -> type[list[Any]]: ...  # pyrefly: ignore[bad-override]
+    @__class__.setter
+    def __class__(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, value: type[list[Any]], /
+    ) -> None: ...
+    def __iter__(self) -> Iterator[_T_co]: ...
+    # copy() is only TEMPORARILY needed because `__class__` is a property
+    # and ty doesn't support property protocol members. Remove when
+    # https://github.com/astral-sh/ty/issues/1379 is resolved
+    def copy(self) -> list[Any]: ...
+
 class SupportsTrueDiv(Protocol[_T_contra, _T_co]):
     def __truediv__(self, x: _T_contra, /) -> _T_co: ...
 

pandas-stubs/core/frame.pyi

@@ -99,6 +99,7 @@ from pandas._typing import (
     CalculationMethod,
     ColspaceArgType,
     CompressionOptions,
+    CovariantList,
     DropKeep,
     Dtype,
     FilePath,
@@ -2092,9 +2093,11 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @overload
     def get(self, key: Hashable, default: _T) -> Series | _T: ...
     @overload
-    def get(self, key: list[Hashable], default: None = None) -> Self | None: ...
+    def get(
+        self, key: CovariantList[Hashable], default: None = None
+    ) -> Self | None: ...
     @overload
-    def get(self, key: list[Hashable], default: _T) -> Self | _T: ...
+    def get(self, key: CovariantList[Hashable], default: _T) -> Self | _T: ...
     def gt(
         self,
         other: complex | ListLike | DataFrame,

tests/frame/test_frame.py

@@ -3921,6 +3921,9 @@ def test_get() -> None:
     )
     check(assert_type(df.get(["z"], default=1), pd.DataFrame | int), int)
 
+    key = ["a", "b"]
+    check(assert_type(df.get(key), pd.DataFrame | None), pd.DataFrame)
+
 
 def test_info() -> None:
     df = pd.DataFrame()

tests/test_typing.py

@@ -1,5 +1,6 @@
 from collections.abc import Mapping
 from typing import (
+    TYPE_CHECKING,
     Any,
     Literal,
 )
@@ -10,9 +11,15 @@
 
 from pandas.core.dtypes.base import ExtensionDtype
 
-from tests import get_dtype
+from tests import (
+    TYPE_CHECKING_INVALID_USAGE,
+    get_dtype,
+)
 from tests.dtypes import DTYPE_ARG_ALIAS_MAPS
 
+if TYPE_CHECKING:
+    from pandas._typing import CovariantList
+
 
 def test_get_dtype() -> None:
     alias_union = (
@@ -36,3 +43,22 @@ def test_dtype_arg_aliases(dtype_arg: Any, alias_map: Mapping[Any, Any]) -> None
     assert set(get_dtype(dtype_arg)) == {
         type(t) if isinstance(t, ExtensionDtype) else t for t in alias_map
     }
+
+
+def test_covariant_list() -> None:
+    def f(_: "CovariantList[float]") -> None: ...
+
+    good1: list[float] = [1.0, 2.0, 3.0]  # OK, trivial case
+    good2: list[int] = [1, 2, 3]  # OK, list[int] < list[float] due to covariance
+    bad1: tuple[float, ...] = (1.0, 2.0, 3.0)  # Error, tuple is not a subtype of list
+    bad2: list[str] = ["a", "b", "c"]  # Error, list[str] !< list[float]
+    bad3: list[object] = [1, "a", 3.0]  # Error, list[object] !< list[float]
+    bad4: float = 1.0  # Error, float !< list[float]
+
+    f(good1)
+    f(good2)
+    if TYPE_CHECKING_INVALID_USAGE:
+        f(bad1)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
+        f(bad2)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
+        f(bad3)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
+        f(bad4)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]