From 9512ad9037630099409b5a161d67fcc5cd3d8786 Mon Sep 17 00:00:00 2001 From: Ran Benita Date: Mon, 27 Sep 2021 12:21:36 +0300 Subject: [PATCH] Fix narrowing of nested union of TypedDicts The code in refine_parent_types() bailed out when it encountered something other than a TypedDict in the union, including a nested union of TypedDicts. This caused the narrowing to fail. Fix it by making it iterate over the flattened union items. Fixes #9308. --- mypy/checker.py | 3 +-- test-data/unit/check-narrowing.test | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 11499d6b570e..047fd2a889ea 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -4676,8 +4676,7 @@ def replay_lookup(new_parent_type: ProperType) -> Optional[Type]: # Take each element in the parent union and replay the original lookup procedure # to figure out which parents are compatible. new_parent_types = [] - for item in parent_type.items: - item = get_proper_type(item) + for item in union_items(parent_type): member_type = replay_lookup(item) if member_type is None: # We were unable to obtain the member type. So, we give up on refining this diff --git a/test-data/unit/check-narrowing.test b/test-data/unit/check-narrowing.test index 7c8415b75fe1..527c89a71b33 100644 --- a/test-data/unit/check-narrowing.test +++ b/test-data/unit/check-narrowing.test @@ -1093,3 +1093,32 @@ def f(t: Type[T], a: A, b: B) -> None: reveal_type(b) # N: Revealed type is "" else: reveal_type(b) # N: Revealed type is "__main__.B" + +[case testNarrowingNestedUnionOfTypedDicts] +from typing import Union +from typing_extensions import Literal, TypedDict + +class A(TypedDict): + tag: Literal["A"] + a: int + +class B(TypedDict): + tag: Literal["B"] + b: int + +class C(TypedDict): + tag: Literal["C"] + c: int + +AB = Union[A, B] +ABC = Union[AB, C] +abc: ABC + +if abc["tag"] == "A": + reveal_type(abc) # N: Revealed type is "TypedDict('__main__.A', {'tag': Literal['A'], 'a': builtins.int})" +elif abc["tag"] == "C": + reveal_type(abc) # N: Revealed type is "TypedDict('__main__.C', {'tag': Literal['C'], 'c': builtins.int})" +else: + reveal_type(abc) # N: Revealed type is "TypedDict('__main__.B', {'tag': Literal['B'], 'b': builtins.int})" + +[builtins fixtures/primitives.pyi]