Move this logic closer to binder, stop putting AssignmentExpr directly

Author: sterliakov

mypy/binder.py

@@ -8,7 +8,16 @@
 
 from mypy.erasetype import remove_instance_last_known_values
 from mypy.literals import Key, extract_var_from_literal_hash, literal, literal_hash, subkeys
-from mypy.nodes import Expression, IndexExpr, MemberExpr, NameExpr, RefExpr, TypeInfo, Var
+from mypy.nodes import (
+    LITERAL_NO,
+    Expression,
+    IndexExpr,
+    MemberExpr,
+    NameExpr,
+    RefExpr,
+    TypeInfo,
+    Var,
+)
 from mypy.options import Options
 from mypy.subtypes import is_same_type, is_subtype
 from mypy.typeops import make_simplified_union
@@ -173,6 +182,15 @@ def _get(self, key: Key, index: int = -1) -> CurrentType | None:
                 return self.frames[i].types[key]
         return None
 
+    @classmethod
+    def can_put_directly(cls, expr: Expression) -> bool:
+        """Will `.put()` on this expression be successful?
+
+        This is inlined in `.put()` because the logic is rather hot and must be kept
+        in sync.
+        """
+        return isinstance(expr, (IndexExpr, MemberExpr, NameExpr)) and literal(expr) > LITERAL_NO
+
     def put(self, expr: Expression, typ: Type, *, from_assignment: bool = True) -> None:
         """Directly set the narrowed type of expression (if it supports it).
 

mypy/checker.py

@@ -5478,6 +5478,9 @@ def visit_continue_stmt(self, s: ContinueStmt) -> None:
 
     def visit_match_stmt(self, s: MatchStmt) -> None:
         named_subject = self._make_named_statement_for_match(s)
+        # In sync with similar actions elsewhere, narrow the target if
+        # we are matching an AssignmentExpr
+        unwrapped_subject = collapse_walrus(s.subject)
         with self.binder.frame_context(can_skip=False, fall_through=0):
             subject_type = get_proper_type(self.expr_checker.accept(s.subject))
 
@@ -5513,9 +5516,9 @@ def visit_match_stmt(self, s: MatchStmt) -> None:
                         # Maybe the subject type can be inferred from constraints on
                         # its attribute/item?
                         if pattern_map and named_subject in pattern_map:
-                            pattern_map[s.subject] = pattern_map[named_subject]
+                            pattern_map[unwrapped_subject] = pattern_map[named_subject]
                         if else_map and named_subject in else_map:
-                            else_map[s.subject] = else_map[named_subject]
+                            else_map[unwrapped_subject] = else_map[named_subject]
                         pattern_map = self.propagate_up_typemap_info(pattern_map)
                         else_map = self.propagate_up_typemap_info(else_map)
                         self.remove_capture_conflicts(pattern_type.captures, inferred_types)
@@ -5571,7 +5574,7 @@ def visit_match_stmt(self, s: MatchStmt) -> None:
     def _make_named_statement_for_match(self, s: MatchStmt) -> Expression:
         """Construct a fake NameExpr for inference if a match clause is complex."""
         subject = s.subject
-        if isinstance(subject, (NameExpr, AssignmentExpr)):
+        if self.binder.can_put_directly(subject):
             # Already named - we should infer type of it as given
             return subject
         elif s.subject_dummy is not None:

test-data/unit/check-python310.test

@@ -2179,9 +2179,11 @@ def f() -> None:
     match x := returns_a_or_none():
         case A():
             reveal_type(x.a)  # N: Revealed type is "builtins.int"
+    reveal_type(x)  # N: Revealed type is "Union[__main__.A, None]"
     match x := returns_a():
         case A():
             reveal_type(x.a)  # N: Revealed type is "builtins.int"
+    reveal_type(x)  # N: Revealed type is "__main__.A"
     y = returns_a_or_none()
     match y:
         case A():