'in' can narrow TypedDict unions

Author: ikonst

mypy/checker.py

@@ -5006,6 +5006,18 @@ def conditional_callable_type_map(
 
         return None, {}
 
+    def _string_in_type_map(self, item_str: str, collection_expr: Expression, t: Type) -> tuple[TypeMap, TypeMap]:
+        if_map, else_map = {}, {}
+        t = get_proper_type(t)
+        if isinstance(t, TypedDictType):
+            (if_map if item_str in t.items else else_map)[collection_expr] = t
+        elif isinstance(t, UnionType):
+            for union_item in t.items:
+                union_if_map, union_else_map = self._string_in_type_map(item_str, collection_expr, union_item)
+                if_map.update(union_if_map)
+                else_map.update(union_else_map)
+        return if_map, else_map
+
     def _is_truthy_type(self, t: ProperType) -> bool:
         return (
             (
@@ -5313,28 +5325,30 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                 elif operator in {"in", "not in"}:
                     assert len(expr_indices) == 2
                     left_index, right_index = expr_indices
-                    if left_index not in narrowable_operand_index_to_hash:
-                        continue
-
                     item_type = operand_types[left_index]
                     collection_type = operand_types[right_index]
 
-                    # We only try and narrow away 'None' for now
-                    if not is_optional(item_type):
-                        continue
+                    if left_index in narrowable_operand_index_to_hash:
+                        # We only try and narrow away 'None' for now
+                        if not is_optional(item_type):
+                            continue
+
+                        collection_item_type = get_proper_type(builtin_item_type(collection_type))
+                        if collection_item_type is None or is_optional(collection_item_type):
+                            continue
+                        if (
+                            isinstance(collection_item_type, Instance)
+                            and collection_item_type.type.fullname == "builtins.object"
+                        ):
+                            continue
+                        if is_overlapping_erased_types(item_type, collection_item_type):
+                            if_map, else_map = {operands[left_index]: remove_optional(item_type)}, {}
+                        else:
+                            continue
+
+                    elif isinstance(item_type.last_known_value, LiteralType) and isinstance(item_type.last_known_value.value, str):
+                        if_map, else_map = self._string_in_type_map(item_type.last_known_value.value, operands[right_index], collection_type)
 
-                    collection_item_type = get_proper_type(builtin_item_type(collection_type))
-                    if collection_item_type is None or is_optional(collection_item_type):
-                        continue
-                    if (
-                        isinstance(collection_item_type, Instance)
-                        and collection_item_type.type.fullname == "builtins.object"
-                    ):
-                        continue
-                    if is_overlapping_erased_types(item_type, collection_item_type):
-                        if_map, else_map = {operands[left_index]: remove_optional(item_type)}, {}
-                    else:
-                        continue
                 else:
                     if_map = {}
                     else_map = {}
@@ -5387,6 +5401,10 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                 or_conditional_maps(left_if_vars, right_if_vars),
                 and_conditional_maps(left_else_vars, right_else_vars),
             )
+        elif isinstance(node, OpExpr) and node.op == "in":
+            left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
+            right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
+
         elif isinstance(node, UnaryExpr) and node.op == "not":
             left, right = self.find_isinstance_check(node.expr)
             return right, left

test-data/unit/check-typeddict.test

@@ -2012,6 +2012,30 @@ v = {bad2: 2}  # E: Extra key "bad" for TypedDict "Value"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
+[case testFinalTypedDictTagged]
+from __future__ import annotations
+from typing import TypedDict
+from typing_extensions import final
+
+@final
+class D1(TypedDict):
+    foo: int
+
+
+@final
+class D2(TypedDict):
+    bar: int
+
+d: D1 | D2
+val = d['foo']  # E: TypedDict "D2" has no key "foo"
+if 'foo' in d:
+    val = d['foo']
+else:
+    val = d['bar']
+
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
 [case testCannotSubclassFinalTypedDict]
 from typing import TypedDict
 from typing_extensions import final