From 9119b14cbcb1cb62fb576c845c84269b0e7b2bcc Mon Sep 17 00:00:00 2001
From: hauntsaninja <>
Date: Tue, 7 Jan 2020 16:43:53 -0800
Subject: [PATCH 1/2] Narrow types for walrus assignment in if statements

---
 mypy/checker.py                    |  2 ++
 test-data/unit/check-python38.test | 12 +++++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 5046d9431b4f..3d286e31800e 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -3937,6 +3937,8 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
                 partial_type_maps.append((if_map, else_map))
 
             return reduce_partial_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/test-data/unit/check-python38.test b/test-data/unit/check-python38.test
index fd0e46ebd1f4..99a732693a10 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'
@@ -294,4 +290,10 @@ def check_partial_list() -> None:
 
     reveal_type(x)  # N: Revealed type is 'builtins.list[builtins.int]'
 
+def check_narrow(x: Optional[int]) -> None:
+    if (y := x):
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+[builtins fixtures/f_string.pyi]
+
+
 [builtins fixtures/f_string.pyi]

From 25b4295774997c8912cf1db31c36d766fe22e677 Mon Sep 17 00:00:00 2001
From: hauntsaninja <>
Date: Tue, 7 Jan 2020 17:48:19 -0800
Subject: [PATCH 2/2] Fix PartialType handling for walrus assignments

---
 mypy/checkexpr.py                  |  1 +
 test-data/unit/check-python38.test | 21 ++++++++++++++-------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index c8dc34da0b91..3374b8174318 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -2772,6 +2772,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 99a732693a10..a4388aeb0299 100644
--- a/test-data/unit/check-python38.test
+++ b/test-data/unit/check-python38.test
@@ -284,16 +284,23 @@ def check_partial() -> None:
 
     reveal_type(x)  # N: Revealed type is 'Union[builtins.int, None]'
 
-def check_partial_list() -> None:
-    if (x := []):
-        x.append(3)
-
-    reveal_type(x)  # N: Revealed type is 'builtins.list[builtins.int]'
-
 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
 
-[builtins fixtures/f_string.pyi]
+def check_partial_list() -> None:
+    if (x := []):  # E: Need type annotation for 'x' (hint: "x: List[<type>] = ...")
+        pass
+
+    y: List[str]
+    if (y := []):
+        pass
+
+    if (z := []):
+        z.append(3)
+    reveal_type(z)  # N: Revealed type is 'builtins.list[builtins.int]'
+[builtins fixtures/list.pyi]