diff --git a/mypy/literals.py b/mypy/literals.py index 6d3a8f2843fe..e325be0ff5fb 100644 --- a/mypy/literals.py +++ b/mypy/literals.py @@ -48,6 +48,7 @@ TypeVarExpr, TypeVarTupleExpr, UnaryExpr, + Var, YieldExpr, YieldFromExpr, ) @@ -112,6 +113,8 @@ def literal(e: Expression) -> int: return LITERAL_NO elif isinstance(e, NameExpr): + if isinstance(e.node, Var) and e.node.is_final and e.node.final_value is not None: + return LITERAL_YES return LITERAL_TYPE if isinstance(e, (IntExpr, FloatExpr, ComplexExpr, StrExpr, BytesExpr)): @@ -154,6 +157,8 @@ def visit_star_expr(self, e: StarExpr) -> Key: return ("Star", literal_hash(e.expr)) def visit_name_expr(self, e: NameExpr) -> Key: + if isinstance(e.node, Var) and e.node.is_final and e.node.final_value is not None: + return ("Literal", e.node.final_value) # N.B: We use the node itself as the key, and not the name, # because using the name causes issues when there is shadowing # (for example, in list comprehensions). diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test index 5046bd3742b5..b6eae1da7d84 100644 --- a/test-data/unit/check-literal.test +++ b/test-data/unit/check-literal.test @@ -1830,7 +1830,8 @@ reveal_type(unify(func)) # N: Revealed type is "" -- [case testLiteralIntelligentIndexingTuples] -from typing import Tuple, NamedTuple +# flags: --strict-optional +from typing import Tuple, NamedTuple, Optional, Final from typing_extensions import Literal class A: pass @@ -1846,17 +1847,23 @@ idx3: Literal[3] idx4: Literal[4] idx5: Literal[5] idx_neg1: Literal[-1] +idx_final: Final = 2 -tup1: Tuple[A, B, C, D, E] +tup1: Tuple[A, B, Optional[C], D, E] reveal_type(tup1[idx0]) # N: Revealed type is "__main__.A" reveal_type(tup1[idx1]) # N: Revealed type is "__main__.B" -reveal_type(tup1[idx2]) # N: Revealed type is "__main__.C" +reveal_type(tup1[idx2]) # N: Revealed type is "Union[__main__.C, None]" +reveal_type(tup1[idx_final]) # N: Revealed type is "Union[__main__.C, None]" reveal_type(tup1[idx3]) # N: Revealed type is "__main__.D" reveal_type(tup1[idx4]) # N: Revealed type is "__main__.E" reveal_type(tup1[idx_neg1]) # N: Revealed type is "__main__.E" tup1[idx5] # E: Tuple index out of range -reveal_type(tup1[idx2:idx4]) # N: Revealed type is "Tuple[__main__.C, __main__.D]" -reveal_type(tup1[::idx2]) # N: Revealed type is "Tuple[__main__.A, __main__.C, __main__.E]" +reveal_type(tup1[idx2:idx4]) # N: Revealed type is "Tuple[Union[__main__.C, None], __main__.D]" +reveal_type(tup1[::idx2]) # N: Revealed type is "Tuple[__main__.A, Union[__main__.C, None], __main__.E]" +if tup1[idx2] is not None: + reveal_type(tup1[idx2]) # N: Revealed type is "Union[__main__.C, None]" +if tup1[idx_final] is not None: + reveal_type(tup1[idx_final]) # N: Revealed type is "__main__.C" Tup2Class = NamedTuple('Tup2Class', [('a', A), ('b', B), ('c', C), ('d', D), ('e', E)]) tup2: Tup2Class @@ -1870,7 +1877,6 @@ tup2[idx5] # E: Tuple index out of range reveal_type(tup2[idx2:idx4]) # N: Revealed type is "Tuple[__main__.C, __main__.D, fallback=__main__.Tup2Class]" reveal_type(tup2[::idx2]) # N: Revealed type is "Tuple[__main__.A, __main__.C, __main__.E, fallback=__main__.Tup2Class]" [builtins fixtures/slice.pyi] -[out] [case testLiteralIntelligentIndexingTypedDict] from typing_extensions import Literal diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 62ac5e31da45..bbde1fad5f29 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -2025,6 +2025,23 @@ class DummyTypedDict(TypedDict): [builtins fixtures/dict.pyi] [typing fixtures/typing-typeddict.pyi] +[case testTypedDictTypeNarrowingWithFinalKey] +from typing import Final, Optional, TypedDict + +KEY_NAME: Final = "bar" +class Foo(TypedDict): + bar: Optional[str] + +foo = Foo(bar="hello") +if foo["bar"] is not None: + reveal_type(foo["bar"]) # N: Revealed type is "builtins.str" + reveal_type(foo[KEY_NAME]) # N: Revealed type is "builtins.str" +if foo[KEY_NAME] is not None: + reveal_type(foo["bar"]) # N: Revealed type is "builtins.str" + reveal_type(foo[KEY_NAME]) # N: Revealed type is "builtins.str" +[builtins fixtures/dict.pyi] +[typing fixtures/typing-typeddict.pyi] + [case testTypedDictDoubleForwardClass] from mypy_extensions import TypedDict from typing import Any, List