Fix type narrowing on TypedDict with key name in Final variable (#11813)

97littleleaf11 · web-flow · commit c8a228989016 · 2022-08-14T20:01:06.000+08:00
* Add support for Final in literal_hash()
* Add test for tuple
diff --git 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
@@ -1830,7 +1830,8 @@ reveal_type(unify(func))  # N: Revealed type is "<nothing>"
 --
 
 [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
@@ -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
