Allow objects matching SupportsKeysAndGetItem to be unpacked (#14990)

bryanforbes · web-flow · commit 14493660eadf · 2023-04-15T13:47:33.000-07:00
Fixes #14986 This PR allows any object matching `_typeshed.SupportsKeysAndGetItem[str, Any]` to be unpacked with `**`.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -2126,7 +2126,9 @@ def check_argument_types(
                 if actual_kind == nodes.ARG_STAR2 and not self.is_valid_keyword_var_arg(
                     actual_type
                 ):
-                    is_mapping = is_subtype(actual_type, self.chk.named_type("typing.Mapping"))
+                    is_mapping = is_subtype(
+                        actual_type, self.chk.named_type("_typeshed.SupportsKeysAndGetItem")
+                    )
                     self.msg.invalid_keyword_var_arg(actual_type, is_mapping, context)
                 expanded_actual = mapper.expand_actual_type(
                     actual_type, actual_kind, callee.arg_names[i], callee_arg_kind
@@ -4346,7 +4348,11 @@ def visit_dict_expr(self, e: DictExpr) -> Type:
             for arg in stargs:
                 if rv is None:
                     constructor = CallableType(
-                        [self.chk.named_generic_type("typing.Mapping", [kt, vt])],
+                        [
+                            self.chk.named_generic_type(
+                                "_typeshed.SupportsKeysAndGetItem", [kt, vt]
+                            )
+                        ],
                         [nodes.ARG_POS],
                         [None],
                         self.chk.named_generic_type("builtins.dict", [kt, vt]),
@@ -4936,14 +4942,14 @@ def is_valid_keyword_var_arg(self, typ: Type) -> bool:
             is_subtype(
                 typ,
                 self.chk.named_generic_type(
-                    "typing.Mapping",
+                    "_typeshed.SupportsKeysAndGetItem",
                     [self.named_type("builtins.str"), AnyType(TypeOfAny.special_form)],
                 ),
             )
             or is_subtype(
                 typ,
                 self.chk.named_generic_type(
-                    "typing.Mapping", [UninhabitedType(), UninhabitedType()]
+                    "_typeshed.SupportsKeysAndGetItem", [UninhabitedType(), UninhabitedType()]
                 ),
             )
             or isinstance(typ, ParamSpecType)
diff --git a/mypy/checkstrformat.py b/mypy/checkstrformat.py
@@ -844,10 +844,14 @@ def build_dict_type(self, expr: FormatStringExpr) -> Type:
         any_type = AnyType(TypeOfAny.special_form)
         if isinstance(expr, BytesExpr):
             bytes_type = self.chk.named_generic_type("builtins.bytes", [])
-            return self.chk.named_generic_type("typing.Mapping", [bytes_type, any_type])
+            return self.chk.named_generic_type(
+                "_typeshed.SupportsKeysAndGetItem", [bytes_type, any_type]
+            )
         elif isinstance(expr, StrExpr):
             str_type = self.chk.named_generic_type("builtins.str", [])
-            return self.chk.named_generic_type("typing.Mapping", [str_type, any_type])
+            return self.chk.named_generic_type(
+                "_typeshed.SupportsKeysAndGetItem", [str_type, any_type]
+            )
         else:
             assert False, "Unreachable"
 
diff --git a/mypy/test/data.py b/mypy/test/data.py
@@ -105,6 +105,12 @@ def parse_test_case(case: DataDrivenTestCase) -> None:
             src_path = join(os.path.dirname(case.file), item.arg)
             with open(src_path, encoding="utf8") as f:
                 files.append((join(base_path, "typing.pyi"), f.read()))
+        elif item.id == "_typeshed":
+            # Use an alternative stub file for the _typeshed module.
+            assert item.arg is not None
+            src_path = join(os.path.dirname(case.file), item.arg)
+            with open(src_path, encoding="utf8") as f:
+                files.append((join(base_path, "_typeshed.pyi"), f.read()))
         elif re.match(r"stale[0-9]*$", item.id):
             passnum = 1 if item.id == "stale" else int(item.id[len("stale") :])
             assert passnum > 0
diff --git a/mypy/test/testdeps.py b/mypy/test/testdeps.py
@@ -50,7 +50,7 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
             type_state.add_all_protocol_deps(deps)
 
             for source, targets in sorted(deps.items()):
-                if source.startswith(("<enum", "<typing", "<mypy")):
+                if source.startswith(("<enum", "<typing", "<mypy", "<_typeshed.")):
                     # Remove noise.
                     continue
                 line = f"{source} -> {', '.join(sorted(targets))}"
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -1,6 +1,7 @@
 # These builtins stubs are used implicitly in AST to IR generation
 # test cases.
 
+import _typeshed
 from typing import (
     TypeVar, Generic, List, Iterator, Iterable, Dict, Optional, Tuple, Any, Set,
     overload, Mapping, Union, Callable, Sequence, FrozenSet, Protocol
diff --git a/mypyc/test-data/fixtures/typing-full.pyi b/mypyc/test-data/fixtures/typing-full.pyi
@@ -125,6 +125,7 @@ class Sequence(Iterable[T_co], Container[T_co]):
     def __getitem__(self, n: Any) -> T_co: pass
 
 class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta):
+    def keys(self) -> Iterable[T]: pass  # Approximate return type
     def __getitem__(self, key: T) -> T_co: pass
     @overload
     def get(self, k: T) -> Optional[T_co]: pass
diff --git a/test-data/unit/check-expressions.test b/test-data/unit/check-expressions.test
@@ -1786,13 +1786,42 @@ b = {'z': 26, *a}  # E: invalid syntax
 
 [case testDictWithStarStarExpr]
 
-from typing import Dict
+from typing import Dict, Iterable
+
+class Thing:
+    def keys(self) -> Iterable[str]:
+        ...
+    def __getitem__(self, key: str) -> int:
+        ...
+
 a = {'a': 1}
 b = {'z': 26, **a}
 c = {**b}
 d = {**a, **b, 'c': 3}
-e = {1: 'a', **a}  # E: Argument 1 to "update" of "dict" has incompatible type "Dict[str, int]"; expected "Mapping[int, str]"
-f = {**b}  # type: Dict[int, int]  # E: List item 0 has incompatible type "Dict[str, int]"; expected "Mapping[int, int]"
+e = {1: 'a', **a}  # E: Argument 1 to "update" of "dict" has incompatible type "Dict[str, int]"; expected "SupportsKeysAndGetItem[int, str]"
+f = {**b}  # type: Dict[int, int]  # E: List item 0 has incompatible type "Dict[str, int]"; expected "SupportsKeysAndGetItem[int, int]"
+g = {**Thing()}
+h = {**a, **Thing()}
+i = {**Thing()}  # type: Dict[int, int]  # E: List item 0 has incompatible type "Thing"; expected "SupportsKeysAndGetItem[int, int]" \
+                 # N: Following member(s) of "Thing" have conflicts: \
+                 # N:     Expected: \
+                 # N:         def __getitem__(self, int, /) -> int \
+                 # N:     Got: \
+                 # N:         def __getitem__(self, str, /) -> int \
+                 # N:     Expected: \
+                 # N:         def keys(self) -> Iterable[int] \
+                 # N:     Got: \
+                 # N:         def keys(self) -> Iterable[str]
+j = {1: 'a', **Thing()}  # E: Argument 1 to "update" of "dict" has incompatible type "Thing"; expected "SupportsKeysAndGetItem[int, str]" \
+                         # N: Following member(s) of "Thing" have conflicts: \
+                         # N:     Expected: \
+                         # N:         def __getitem__(self, int, /) -> str \
+                         # N:     Got: \
+                         # N:         def __getitem__(self, str, /) -> int \
+                         # N:     Expected: \
+                         # N:         def keys(self) -> Iterable[int] \
+                         # N:     Got: \
+                         # N:         def keys(self) -> Iterable[str]
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-medium.pyi]
 
diff --git a/test-data/unit/check-formatting.test b/test-data/unit/check-formatting.test
@@ -125,14 +125,29 @@ b'%(x)s' % {b'x': b'data'}
 [typing fixtures/typing-medium.pyi]
 
 [case testStringInterpolationMappingDictTypes]
-from typing import Any, Dict
+from typing import Any, Dict, Iterable
+
+class StringThing:
+    def keys(self) -> Iterable[str]:
+        ...
+    def __getitem__(self, __key: str) -> str:
+        ...
+
+class BytesThing:
+    def keys(self) -> Iterable[bytes]:
+        ...
+    def __getitem__(self, __key: bytes) -> str:
+        ...
+
 a = None # type: Any
 ds, do, di = None, None, None # type: Dict[str, int], Dict[object, int], Dict[int, int]
-'%(a)' % 1  # E: Format requires a mapping (expression has type "int", expected type for mapping is "Mapping[str, Any]")
+'%(a)' % 1  # E: Format requires a mapping (expression has type "int", expected type for mapping is "SupportsKeysAndGetItem[str, Any]")
 '%()d' % a
 '%()d' % ds
-'%()d' % do  # E: Format requires a mapping (expression has type "Dict[object, int]", expected type for mapping is "Mapping[str, Any]")
-b'%()d' % ds  # E: Format requires a mapping (expression has type "Dict[str, int]", expected type for mapping is "Mapping[bytes, Any]")
+'%()d' % do  # E: Format requires a mapping (expression has type "Dict[object, int]", expected type for mapping is "SupportsKeysAndGetItem[str, Any]")
+b'%()d' % ds  # E: Format requires a mapping (expression has type "Dict[str, int]", expected type for mapping is "SupportsKeysAndGetItem[bytes, Any]")
+'%()s' % StringThing()
+b'%()s' % BytesThing()
 [builtins fixtures/primitives.pyi]
 
 [case testStringInterpolationMappingInvalidSpecifiers]
diff --git a/test-data/unit/check-generic-subtyping.test b/test-data/unit/check-generic-subtyping.test
@@ -990,6 +990,7 @@ main:13: note: Revealed type is "builtins.dict[builtins.int, builtins.str]"
 main:14: error: Keywords must be strings
 main:14: error: Argument 1 to "func_with_kwargs" has incompatible type "**X1[str, int]"; expected "int"
 [builtins fixtures/dict.pyi]
+[typing fixtures/typing-medium.pyi]
 
 [case testSubtypingMappingUnpacking3]
 from typing import Generic, TypeVar, Mapping, Iterable
diff --git a/test-data/unit/check-incremental.test b/test-data/unit/check-incremental.test
@@ -3699,8 +3699,8 @@ cache_fine_grained = False
 [file mypy.ini.2]
 \[mypy]
 cache_fine_grained = True
-[rechecked a, builtins, typing]
-[stale a, builtins, typing]
+[rechecked _typeshed, a, builtins, typing]
+[stale _typeshed, a, builtins, typing]
 [builtins fixtures/tuple.pyi]
 
 [case testIncrementalPackageNameOverload]
@@ -3751,8 +3751,8 @@ Signature: 8a477f597d28d172789f06886806bc55
 [file b.py.2]
 # uh
 -- Every file should get reloaded, since the cache was invalidated
-[stale a, b, builtins, typing]
-[rechecked a, b, builtins, typing]
+[stale _typeshed, a, b, builtins, typing]
+[rechecked _typeshed, a, b, builtins, typing]
 [builtins fixtures/tuple.pyi]
 
 [case testIncrementalBustedFineGrainedCache2]
@@ -3764,8 +3764,8 @@ import b
 [file b.py.2]
 # uh
 -- Every file should get reloaded, since the settings changed
-[stale a, b, builtins, typing]
-[rechecked a, b, builtins, typing]
+[stale _typeshed, a, b, builtins, typing]
+[rechecked _typeshed, a, b, builtins, typing]
 [builtins fixtures/tuple.pyi]
 
 [case testIncrementalBustedFineGrainedCache3]
@@ -3780,8 +3780,8 @@ import b
 [file b.py.2]
 # uh
 -- Every file should get reloaded, since the cache was invalidated
-[stale a, b, builtins, typing]
-[rechecked a, b, builtins, typing]
+[stale _typeshed, a, b, builtins, typing]
+[rechecked _typeshed, a, b, builtins, typing]
 [builtins fixtures/tuple.pyi]
 
 [case testIncrementalWorkingFineGrainedCache]
diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test
@@ -1671,7 +1671,9 @@ a() # E: "Dict[str, int]" not callable
 
 [case testInferDictInitializedToEmptyUsingUpdateError]
 a = {}  # E: Need type annotation for "a" (hint: "a: Dict[<type>, <type>] = ...")
-a.update([1, 2])  # E: Argument 1 to "update" of "dict" has incompatible type "List[int]"; expected "Mapping[Any, Any]"
+a.update([1, 2])  # E: Argument 1 to "update" of "dict" has incompatible type "List[int]"; expected "SupportsKeysAndGetItem[Any, Any]" \
+                  # N: "list" is missing following "SupportsKeysAndGetItem" protocol member: \
+                  # N:     keys
 a()  # E: "Dict[Any, Any]" not callable
 [builtins fixtures/dict.pyi]
 
diff --git a/test-data/unit/check-kwargs.test b/test-data/unit/check-kwargs.test
@@ -499,7 +499,7 @@ g(**{})
 
 [case testKeywordUnpackWithDifferentTypes]
 # https://github.com/python/mypy/issues/11144
-from typing import Dict, Generic, TypeVar, Mapping
+from typing import Dict, Generic, TypeVar, Mapping, Iterable
 
 T = TypeVar("T")
 T2 = TypeVar("T2")
@@ -516,21 +516,29 @@ class C(Generic[T, T2]):
 class D:
     ...
 
+class E:
+    def keys(self) -> Iterable[str]:
+        ...
+    def __getitem__(self, key: str) -> float:
+        ...
+
 def foo(**i: float) -> float:
     ...
 
 a: A[str, str]
 b: B[str, str]
 c: C[str, float]
 d: D
-e = {"a": "b"}
+e: E
+f = {"a": "b"}
 
 foo(k=1.5)
 foo(**a)
 foo(**b)
 foo(**c)
 foo(**d)
 foo(**e)
+foo(**f)
 
 # Correct:
 
@@ -544,9 +552,9 @@ foo(**good1)
 foo(**good2)
 foo(**good3)
 [out]
-main:29: error: Argument 1 to "foo" has incompatible type "**A[str, str]"; expected "float"
-main:30: error: Argument 1 to "foo" has incompatible type "**B[str, str]"; expected "float"
-main:31: error: Argument after ** must be a mapping, not "C[str, float]"
-main:32: error: Argument after ** must be a mapping, not "D"
-main:33: error: Argument 1 to "foo" has incompatible type "**Dict[str, str]"; expected "float"
+main:36: error: Argument 1 to "foo" has incompatible type "**A[str, str]"; expected "float"
+main:37: error: Argument 1 to "foo" has incompatible type "**B[str, str]"; expected "float"
+main:38: error: Argument after ** must be a mapping, not "C[str, float]"
+main:39: error: Argument after ** must be a mapping, not "D"
+main:41: error: Argument 1 to "foo" has incompatible type "**Dict[str, str]"; expected "float"
 [builtins fixtures/dict.pyi]
diff --git a/test-data/unit/fine-grained-dataclass-transform.test b/test-data/unit/fine-grained-dataclass-transform.test
@@ -86,9 +86,9 @@ class A(Dataclass):
 
 [out]
 main:7: error: Unexpected keyword argument "x" for "B"
-builtins.pyi:12: note: "B" defined here
+builtins.pyi:13: note: "B" defined here
 main:7: error: Unexpected keyword argument "y" for "B"
-builtins.pyi:12: note: "B" defined here
+builtins.pyi:13: note: "B" defined here
 ==
 
 [case frozenInheritanceViaDefault]
diff --git a/test-data/unit/fine-grained-modules.test b/test-data/unit/fine-grained-modules.test
@@ -1279,12 +1279,12 @@ a.py:2: error: Too many arguments for "foo"
 
 [case testAddModuleAfterCache3-only_when_cache]
 # cmd: mypy main a.py
-# cmd2: mypy main a.py b.py c.py d.py e.py f.py g.py h.py
-# cmd3: mypy main a.py b.py c.py d.py e.py f.py g.py h.py
+# cmd2: mypy main a.py b.py c.py d.py e.py f.py g.py h.py i.py j.py
+# cmd3: mypy main a.py b.py c.py d.py e.py f.py g.py h.py i.py j.py
 # flags: --ignore-missing-imports --follow-imports=skip
 import a
 [file a.py]
-import b, c, d, e, f, g, h
+import b, c, d, e, f, g, h, i, j
 b.foo(10)
 [file b.py.2]
 def foo() -> None: pass
@@ -1294,6 +1294,8 @@ def foo() -> None: pass
 [file f.py.2]
 [file g.py.2]
 [file h.py.2]
+[file i.py.2]
+[file j.py.2]
 
 -- No files should be stale or reprocessed in the first step since the large number
 -- of missing files will force build to give up on cache loading.
diff --git a/test-data/unit/fine-grained.test b/test-data/unit/fine-grained.test
@@ -7546,7 +7546,7 @@ def d() -> Dict[int, int]: pass
 [builtins fixtures/dict.pyi]
 [out]
 ==
-main:5: error: Argument 1 to "update" of "dict" has incompatible type "Dict[int, int]"; expected "Mapping[int, str]"
+main:5: error: Argument 1 to "update" of "dict" has incompatible type "Dict[int, int]"; expected "SupportsKeysAndGetItem[int, str]"
 
 [case testAwaitAndAsyncDef-only_when_nocache]
 from a import g
diff --git a/test-data/unit/fixtures/args.pyi b/test-data/unit/fixtures/args.pyi
@@ -1,5 +1,6 @@
 # Builtins stub used to support *args, **kwargs.
 
+import _typeshed
 from typing import TypeVar, Generic, Iterable, Sequence, Tuple, Dict, Any, overload, Mapping
 
 Tco = TypeVar('Tco', covariant=True)
diff --git a/test-data/unit/fixtures/dataclasses.pyi b/test-data/unit/fixtures/dataclasses.pyi
@@ -1,3 +1,4 @@
+import _typeshed
 from typing import (
     Generic, Iterator, Iterable, Mapping, Optional, Sequence, Tuple,
     TypeVar, Union, overload,
diff --git a/test-data/unit/fixtures/dict.pyi b/test-data/unit/fixtures/dict.pyi
@@ -1,5 +1,7 @@
 # Builtins stub used in dictionary-related test cases.
 
+from _typeshed import SupportsKeysAndGetItem
+import _typeshed
 from typing import (
     TypeVar, Generic, Iterable, Iterator, Mapping, Tuple, overload, Optional, Union, Sequence
 )
@@ -25,7 +27,7 @@ class dict(Mapping[KT, VT]):
     def __setitem__(self, k: KT, v: VT) -> None: pass
     def __iter__(self) -> Iterator[KT]: pass
     def __contains__(self, item: object) -> int: pass
-    def update(self, a: Mapping[KT, VT]) -> None: pass
+    def update(self, a: SupportsKeysAndGetItem[KT, VT]) -> None: pass
     @overload
     def get(self, k: KT) -> Optional[VT]: pass
     @overload
diff --git a/test-data/unit/fixtures/paramspec.pyi b/test-data/unit/fixtures/paramspec.pyi
@@ -1,5 +1,6 @@
 # builtins stub for paramspec-related test cases
 
+import _typeshed
 from typing import (
     Sequence, Generic, TypeVar, Iterable, Iterator, Tuple, Mapping, Optional, Union, Type, overload,
     Protocol
diff --git a/test-data/unit/fixtures/primitives.pyi b/test-data/unit/fixtures/primitives.pyi
@@ -1,4 +1,5 @@
 # builtins stub with non-generic primitive types
+import _typeshed
 from typing import Generic, TypeVar, Sequence, Iterator, Mapping, Iterable, Tuple, Union
 
 T = TypeVar('T')
diff --git a/test-data/unit/fixtures/tuple.pyi b/test-data/unit/fixtures/tuple.pyi
@@ -1,5 +1,6 @@
 # Builtins stub used in tuple-related test cases.
 
+import _typeshed
 from typing import Iterable, Iterator, TypeVar, Generic, Sequence, Optional, overload, Tuple, Type
 
 T = TypeVar("T")
diff --git a/test-data/unit/fixtures/typing-async.pyi b/test-data/unit/fixtures/typing-async.pyi
@@ -108,6 +108,7 @@ class Sequence(Iterable[T_co], Container[T_co]):
     def __getitem__(self, n: Any) -> T_co: pass
 
 class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta):
+    def keys(self) -> Iterable[T]: pass  # Approximate return type
     def __getitem__(self, key: T) -> T_co: pass
     @overload
     def get(self, k: T) -> Optional[T_co]: pass
diff --git a/test-data/unit/fixtures/typing-full.pyi b/test-data/unit/fixtures/typing-full.pyi
diff --git a/test-data/unit/fixtures/typing-medium.pyi b/test-data/unit/fixtures/typing-medium.pyi
diff --git a/test-data/unit/fixtures/typing-namedtuple.pyi b/test-data/unit/fixtures/typing-namedtuple.pyi
diff --git a/test-data/unit/fixtures/typing-typeddict.pyi b/test-data/unit/fixtures/typing-typeddict.pyi
diff --git a/test-data/unit/lib-stub/_typeshed.pyi b/test-data/unit/lib-stub/_typeshed.pyi
diff --git a/test-data/unit/lib-stub/builtins.pyi b/test-data/unit/lib-stub/builtins.pyi
diff --git a/test-data/unit/lib-stub/typing.pyi b/test-data/unit/lib-stub/typing.pyi
