feat: proper narrowing for TypedDict keys and values

Author: BobTheBuidler

mypy/checkmember.py

@@ -1360,7 +1360,29 @@ def analyze_enum_class_attribute_access(
 def analyze_typeddict_access(
     name: str, typ: TypedDictType, mx: MemberContext, override_info: TypeInfo | None
 ) -> Type:
-    if name == "__setitem__":
+    if name == "keys":
+        # Return KeysView[union of Literal key types]
+        keys_view_info = mx.chk.named_type("typing.KeysView").type
+        return CallableType(
+            arg_types=[],
+            arg_kinds=[],
+            arg_names=[],
+            ret_type=Instance(keys_view_info, [typ.key_type]),
+            fallback=mx.chk.named_type("builtins.function"),
+            name=name,
+        )
+    elif name == "values":
+        # Return ValuesView[union of value types]
+        values_view_info = mx.chk.named_type("typing.ValuesView").type
+        return CallableType(
+            arg_types=[],
+            arg_kinds=[],
+            arg_names=[],
+            ret_type=Instance(values_view_info, [typ.value_type]),
+            fallback=mx.chk.named_type("builtins.function"),
+            name=name,
+        )
+    elif name == "__setitem__":
         if isinstance(mx.context, IndexExpr):
             # Since we can get this during `a['key'] = ...`
             # it is safe to assume that the context is `IndexExpr`.

mypy/types.py

@@ -2584,6 +2584,18 @@ def __init__(
         self.extra_items_from = []
         self.to_be_mutated = False
 
+    @property
+    def key_type(self) -> Type:
+        """Return a Union of Literal types for all keys."""
+        return UnionType.make_union(
+            [LiteralType(key, self.fallback) for key in self.items.keys()]
+        )
+
+    @property
+    def value_type(self) -> Type:
+        """Return a Union of all value types (deduplicated)."""
+        return UnionType.make_union(list({get_proper_type(typ) for typ in self.items.values()}))
+
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_typeddict_type(self)
 

test-data/unit/check-typeddict.test

@@ -5,8 +5,8 @@ from typing import TypedDict
 Point = TypedDict('Point', {'x': int, 'y': int})
 p = Point(x=42, y=1337)
 reveal_type(p)  # N: Revealed type is "TypedDict('__main__.Point', {'x': builtins.int, 'y': builtins.int})"
-# Use values() to check fallback value type.
-reveal_type(p.values()) # N: Revealed type is "typing.Iterable[builtins.object]"
+reveal_type(p.values()) # N: Revealed type is "typing.ValuesView[builtins.int]"
+reveal_type(p.keys())   # N: Revealed type is "typing.KeysView[Union[Literal['x'], Literal['y']]]"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 [targets __main__]
@@ -16,8 +16,8 @@ from typing import TypedDict
 Point = TypedDict('Point', {'x': int, 'y': int})
 p = Point(dict(x=42, y=1337))
 reveal_type(p)  # N: Revealed type is "TypedDict('__main__.Point', {'x': builtins.int, 'y': builtins.int})"
-# Use values() to check fallback value type.
-reveal_type(p.values()) # N: Revealed type is "typing.Iterable[builtins.object]"
+reveal_type(p.values()) # N: Revealed type is "typing.ValuesView[builtins.int]"
+reveal_type(p.keys())   # N: Revealed type is "typing.KeysView[Union[Literal['x'], Literal['y']]]"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -26,8 +26,8 @@ from typing import TypedDict
 Point = TypedDict('Point', {'x': int, 'y': int})
 p = Point({'x': 42, 'y': 1337})
 reveal_type(p)  # N: Revealed type is "TypedDict('__main__.Point', {'x': builtins.int, 'y': builtins.int})"
-# Use values() to check fallback value type.
-reveal_type(p.values()) # N: Revealed type is "typing.Iterable[builtins.object]"
+reveal_type(p.values()) # N: Revealed type is "typing.ValuesView[builtins.int]"
+reveal_type(p.keys())   # N: Revealed type is "typing.KeysView[Union[Literal['x'], Literal['y']]]"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -36,7 +36,8 @@ from typing import TypedDict, TypeVar, Union
 EmptyDict = TypedDict('EmptyDict', {})
 p = EmptyDict()
 reveal_type(p)  # N: Revealed type is "TypedDict('__main__.EmptyDict', {})"
-reveal_type(p.values()) # N: Revealed type is "typing.Iterable[builtins.object]"
+reveal_type(p.values()) # N: Revealed type is "typing.ValuesView[Never]"
+reveal_type(p.keys())   # N: Revealed type is "typing.KeysView[Never]"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -534,8 +535,8 @@ Point3D = TypedDict('Point3D', {'x': int, 'y': int, 'z': int})
 p1 = TaggedPoint(type='2d', x=0, y=0)
 p2 = Point3D(x=1, y=1, z=1)
 joined_points = [p1, p2][0]
-reveal_type(p1.values())   # N: Revealed type is "typing.Iterable[builtins.object]"
-reveal_type(p2.values())   # N: Revealed type is "typing.Iterable[builtins.object]"
+reveal_type(p1.values())   # N: Revealed type is "typing.ValuesView[Union[builtins.str, builtins.int]]"
+reveal_type(p2.values())   # N: Revealed type is "typing.ValuesView[builtins.int]"
 reveal_type(joined_points)  # N: Revealed type is "TypedDict({'x': builtins.int, 'y': builtins.int})"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]

test-data/unit/fixtures/typing-typeddict.pyi

@@ -79,3 +79,7 @@ class _TypedDict(Mapping[str, object]):
     def __delitem__(self, k: NoReturn) -> None: ...
 
 class _SpecialForm: pass
+
+class KeysView(Iterable[T]): pass
+
+class ValuesView(Iterable[V]): pass