From fd08d98dbe8cbbc755c4392654acdf71a497dcf6 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Tue, 4 May 2021 16:30:02 +0100 Subject: [PATCH] Fix narrowing down TypedDict unions with enum literal types A union of enum literals was merged back to the enum type, which broke type narrowing. Disable merging in this case to work around the issue. The fix feels a bit ad hoc. However, I'd rather not spend a lot of time figuring out a general fix, since this seems like a pretty rare edge case. Fixes #10414. --- mypy/checkexpr.py | 4 +++- test-data/unit/check-narrowing.test | 23 +++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 1da6150be031..4a64c01dc479 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2910,9 +2910,11 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr, if isinstance(left_type, UnionType): original_type = original_type or left_type + # Don't combine literal types, since we may need them for type narrowing. return make_simplified_union([self.visit_index_with_type(typ, e, original_type) - for typ in left_type.relevant_items()]) + for typ in left_type.relevant_items()], + contract_literals=False) elif isinstance(left_type, TupleType) and self.chk.in_checked_function(): # Special case for tuples. They return a more specific type when # indexed by an integer literal. diff --git a/test-data/unit/check-narrowing.test b/test-data/unit/check-narrowing.test index 14c2e890e122..1697432d3d8b 100644 --- a/test-data/unit/check-narrowing.test +++ b/test-data/unit/check-narrowing.test @@ -1029,3 +1029,26 @@ else: reveal_type(str_or_bool_literal) # N: Revealed type is "Union[Literal[False], Literal[True]]" [builtins fixtures/primitives.pyi] + +[case testNarrowingTypedDictUsingEnumLiteral] +from typing import Union +from typing_extensions import TypedDict, Literal +from enum import Enum + +class E(Enum): + FOO = "a" + BAR = "b" + +class Foo(TypedDict): + tag: Literal[E.FOO] + x: int + +class Bar(TypedDict): + tag: Literal[E.BAR] + y: int + +def f(d: Union[Foo, Bar]) -> None: + assert d['tag'] == E.FOO + d['x'] + reveal_type(d) # N: Revealed type is "TypedDict('__main__.Foo', {'tag': Literal[__main__.E.FOO], 'x': builtins.int})" +[builtins fixtures/dict.pyi]