From 33fc8a3f308adec1929d8bc09bd1860684c71c69 Mon Sep 17 00:00:00 2001
From: hauntsaninja <hauntsaninja@users.noreply.github.com>
Date: Wed, 8 Jan 2020 13:40:54 -0800
Subject: [PATCH] Narrow types for walrus assignment in if statements (#8258)

Fixes #7313
Fixes #7316
---
 mypy/checker.py                    |  2 ++
 mypy/checkexpr.py                  |  1 +
 test-data/unit/check-python38.test | 27 ++++++++++++++++++---------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 18f3573f14c5..9c02d08ebd34 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -3982,6 +3982,8 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                 partial_type_maps.append((if_map, else_map))
 
             return reduce_conditional_maps(partial_type_maps)
+        elif isinstance(node, AssignmentExpr):
+            return self.find_isinstance_check_helper(node.target)
         elif isinstance(node, RefExpr):
             # Restrict the type of the variable to True-ish/False-ish in the if and else branches
             # respectively
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 47f0e74691c0..9a0815060bca 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -2773,6 +2773,7 @@ def visit_assignment_expr(self, e: AssignmentExpr) -> Type:
         value = self.accept(e.value)
         self.chk.check_assignment(e.target, e.value)
         self.chk.check_final(e)
+        self.find_partial_type_ref_fast_path(e.target)
         return value
 
     def visit_unary_expr(self, e: UnaryExpr) -> Type:
diff --git a/test-data/unit/check-python38.test b/test-data/unit/check-python38.test
index fd0e46ebd1f4..a4388aeb0299 100644
--- a/test-data/unit/check-python38.test
+++ b/test-data/unit/check-python38.test
@@ -266,11 +266,7 @@ def check_binder(x: Optional[int], y: Optional[int], z: Optional[int], a: Option
         reveal_type(y)  # N: Revealed type is 'Union[builtins.int, None]'
 
     if x and (y := 1):
-        # TODO should just be int
-        # This is because in check_boolean_op in checkexpr.py we accept the right conditional
-        # within a binder frame context, so the types assigned in it are lost later. Perhaps
-        # we need to make find_isinstance_check() walrus-aware.
-        reveal_type(y)  # N: Revealed type is 'Union[builtins.int, None]'
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
 
     if (a := 1) and x:
         reveal_type(a)  # N: Revealed type is 'builtins.int'
@@ -288,10 +284,23 @@ def check_partial() -> None:
 
     reveal_type(x)  # N: Revealed type is 'Union[builtins.int, None]'
 
+def check_narrow(x: Optional[int]) -> None:
+    if (y := x):
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+[builtins fixtures/f_string.pyi]
+
+[case testWalrusPartialTypes]
+from typing import List
+
 def check_partial_list() -> None:
-    if (x := []):
-        x.append(3)
+    if (x := []):  # E: Need type annotation for 'x' (hint: "x: List[<type>] = ...")
+        pass
 
-    reveal_type(x)  # N: Revealed type is 'builtins.list[builtins.int]'
+    y: List[str]
+    if (y := []):
+        pass
 
-[builtins fixtures/f_string.pyi]
+    if (z := []):
+        z.append(3)
+    reveal_type(z)  # N: Revealed type is 'builtins.list[builtins.int]'
+[builtins fixtures/list.pyi]