Support narrowing of walrus in most cases (#8458)

mypy/checker.py

@@ -3924,10 +3924,12 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
         elif is_false_literal(node):
             return None, {}
         elif isinstance(node, CallExpr):
+            if len(node.args) == 0:
+                return {}, {}
+            expr = collapse_walrus(node.args[0])
             if refers_to_fullname(node.callee, 'builtins.isinstance'):
                 if len(node.args) != 2:  # the error will be reported elsewhere
                     return {}, {}
-                expr = node.args[0]
                 if literal(expr) == LITERAL_TYPE:
                     return self.conditional_type_map_with_intersection(
                         expr,
@@ -3937,13 +3939,11 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
             elif refers_to_fullname(node.callee, 'builtins.issubclass'):
                 if len(node.args) != 2:  # the error will be reported elsewhere
                     return {}, {}
-                expr = node.args[0]
                 if literal(expr) == LITERAL_TYPE:
                     return self.infer_issubclass_maps(node, expr, type_map)
             elif refers_to_fullname(node.callee, 'builtins.callable'):
                 if len(node.args) != 1:  # the error will be reported elsewhere
                     return {}, {}
-                expr = node.args[0]
                 if literal(expr) == LITERAL_TYPE:
                     vartype = type_map[expr]
                     return self.conditional_callable_type_map(expr, vartype)
@@ -3952,7 +3952,7 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
             # narrow their types. (For example, we shouldn't try narrowing the
             # types of literal string or enum expressions).
 
-            operands = node.operands
+            operands = [collapse_walrus(x) for x in node.operands]
             operand_types = []
             narrowable_operand_index_to_hash = {}
             for i, expr in enumerate(operands):
@@ -5742,3 +5742,14 @@ def has_bool_item(typ: ProperType) -> bool:
         return any(is_named_instance(item, 'builtins.bool')
                    for item in typ.items)
     return False
+
+
+def collapse_walrus(e: Expression) -> Expression:
+    """If an expression is an AssignmentExpr, pull out the assignment target.
+
+    We don't make any attempt to pull out all the targets in code like `x := (y := z)`.
+    We could support narrowing those if that sort of code turns out to be common.
+    """
+    if isinstance(e, AssignmentExpr):
+        return e.target
+    return e

test-data/unit/check-python38.test

@@ -189,7 +189,7 @@ def f(p1: bytes, p2: float, /) -> None:
 
 [case testWalrus]
 # flags: --strict-optional
-from typing import NamedTuple, Optional
+from typing import NamedTuple, Optional, List
 from typing_extensions import Final
 
 if a := 2:
@@ -288,10 +288,23 @@ def check_partial() -> None:
 
     reveal_type(x)  # N: Revealed type is 'Union[builtins.int, None]'
 
-def check_narrow(x: Optional[int]) -> None:
+def check_narrow(x: Optional[int], s: List[int]) -> None:
     if (y := x):
         reveal_type(y)  # N: Revealed type is 'builtins.int'
-[builtins fixtures/f_string.pyi]
+
+    if (y := x) is not None:
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+
+    if (y := x) == 10:
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+
+    if (y := x) in s:
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+
+    if isinstance((y := x), int):
+        reveal_type(y)  # N: Revealed type is 'builtins.int'
+
+[builtins fixtures/isinstancelist.pyi]
 
 [case testWalrusPartialTypes]
 from typing import List