Narrow types for walrus assignment in if statements

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

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:

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]