Repository: python/mypy

Issue:
PEP 572: Type not narrowed correctly in assignment
Example from @gvanrossum in #6899:
```
from typing import Optional

class C:
    def foo(self):
        pass

def good(b: Optional[C]) -> None:
    a = b
    if a:
        a.foo()

def bad(b: Optional[C]) -> None:
    if a := b:
        a.foo()  # E:  Item "None" of "Optional[C]" has no attribute "foo"
```


Discussion:
This is actually pretty important to support, this example:
```python
discount = 0.0
if (mo := re.search(r'(\d+)% discount', advertisement)):
    discount = float(mo.group(1)) / 100.0
```
is the flagship example taken from Python 3.8 What's New, and is actually the prime use case for assignment expressions in general. Raising priority to high.

Resolving patch:
diff --git a/mypy/checker.py b/mypy/checker.py
--- 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
--- 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:


Test patch:
diff --git a/test-data/unit/check-python38.test b/test-data/unit/check-python38.test
--- 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]