Narrow types for walrus in identity comparison

In theory, this draft commit fixes python#8236

However, it looks like the ComparisonExpr part of
find_isinstance_check_helper was rewritten today morning in python#8151.
It looks like it's not super easy to merge: `is_valid_target` gets in the
way and I'm a lot less sure about the change.

I'm also not fully sure about the implications of making
AssignmentExpr into mypy.literals, but seems like it's what we'd want.

I also don't really like the way I've branched for AssignmentExpr, it's
pretty unsatisfactory. Well, work in progress!

mypy/checker.py

@@ -4245,6 +4245,8 @@ def refine_identity_comparison_expression(self,
             if enum_name is not None:
                 expr_type = try_expanding_enum_to_union(expr_type, enum_name)
             partial_type_maps.append(conditional_type_map(expr, expr_type, target_type))
+            if isinstance(expr, AssignmentExpr):
+                partial_type_maps.append(conditional_type_map(expr.target, expr_type, target_type))
 
         return reduce_partial_conditional_maps(partial_type_maps)
 

mypy/literals.py

@@ -67,6 +67,9 @@ def literal(e: Expression) -> int:
         else:
             return LITERAL_NO
 
+    elif isinstance(e, AssignmentExpr):
+        return LITERAL_TYPE
+
     elif isinstance(e, NameExpr):
         return LITERAL_TYPE
 
@@ -118,6 +121,9 @@ def visit_name_expr(self, e: NameExpr) -> Key:
         # (for example, in list comprehensions).
         return ('Var', e.node)
 
+    def visit_assignment_expr(self, e: AssignmentExpr) -> None:
+        return ('Assignment', literal_hash(e.target), literal_hash(e.value))
+
     def visit_member_expr(self, e: MemberExpr) -> Key:
         return ('Member', literal_hash(e.expr), e.name)
 
@@ -159,9 +165,6 @@ def visit_index_expr(self, e: IndexExpr) -> Optional[Key]:
             return ('Index', literal_hash(e.base), literal_hash(e.index))
         return None
 
-    def visit_assignment_expr(self, e: AssignmentExpr) -> None:
-        return None
-
     def visit_call_expr(self, e: CallExpr) -> None:
         return None
 

test-data/unit/check-python38.test

@@ -287,6 +287,11 @@ def check_partial() -> None:
 def check_narrow(x: Optional[int]) -> None:
     if (y := x):
         reveal_type(y)  # N: Revealed type is 'builtins.int'
+
+from typing import List
+
+def check_comparison(mystery: List[Optional[str]]) -> None:
+    [a.format() for a in mystery if (a := a) is not None]
 [builtins fixtures/f_string.pyi]
 
 [case testWalrusPartialTypes]