Fix narrow for walrus expressions used in match statements (#14844)

Fixes #14843
python · Mar 6, 2023 · 31f70d7 · 31f70d7
1 parent da6e4a0
commit 31f70d7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -6763,6 +6763,7 @@ def conditional_types(
 def conditional_types_to_typemaps(
     expr: Expression, yes_type: Type | None, no_type: Type | None
 ) -> tuple[TypeMap, TypeMap]:
+    expr = collapse_walrus(expr)
     maps: list[TypeMap] = []
     for typ in (yes_type, no_type):
         proper_type = get_proper_type(typ)

diff --git a/test-data/unit/check-python310.test b/test-data/unit/check-python310.test
@@ -1841,3 +1841,25 @@ class D: pass
 X = None | C
 Y = None | D
 [builtins fixtures/type.pyi]
+
+[case testMatchStatementWalrus]
+class A:
+    a = 1
+
+def returns_a_or_none() -> A | None:
+    return A()
+
+def returns_a() -> A:
+    return A()
+
+def f() -> None:
+    match x := returns_a_or_none():
+        case A():
+            reveal_type(x.a)  # N: Revealed type is "builtins.int"
+    match x := returns_a():
+        case A():
+            reveal_type(x.a)  # N: Revealed type is "builtins.int"
+    y = returns_a_or_none()
+    match y:
+        case A():
+            reveal_type(y.a)  # N: Revealed type is "builtins.int"