Skip to content

Commit 32a0185

Browse files
committed
Support narrowing of walrus in most cases
It is a pretty simple matter of pulling out the assignment target from the walrus. We don't bother handling things like `x := (y := z)` since I can't imagine they are common enough to be worth bothering but we could in the future if anyone cares. Fixes #8447.
1 parent 30c46ab commit 32a0185

File tree

2 files changed

+30
-8
lines changed

2 files changed

+30
-8
lines changed

mypy/checker.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3923,11 +3923,11 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39233923
return {}, None
39243924
elif is_false_literal(node):
39253925
return None, {}
3926-
elif isinstance(node, CallExpr):
3926+
elif isinstance(node, CallExpr) and len(node.args) > 0:
3927+
expr = collapse_walrus(node.args[0])
39273928
if refers_to_fullname(node.callee, 'builtins.isinstance'):
39283929
if len(node.args) != 2: # the error will be reported elsewhere
39293930
return {}, {}
3930-
expr = node.args[0]
39313931
if literal(expr) == LITERAL_TYPE:
39323932
return self.conditional_type_map_with_intersection(
39333933
expr,
@@ -3937,13 +3937,11 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39373937
elif refers_to_fullname(node.callee, 'builtins.issubclass'):
39383938
if len(node.args) != 2: # the error will be reported elsewhere
39393939
return {}, {}
3940-
expr = node.args[0]
39413940
if literal(expr) == LITERAL_TYPE:
39423941
return self.infer_issubclass_maps(node, expr, type_map)
39433942
elif refers_to_fullname(node.callee, 'builtins.callable'):
39443943
if len(node.args) != 1: # the error will be reported elsewhere
39453944
return {}, {}
3946-
expr = node.args[0]
39473945
if literal(expr) == LITERAL_TYPE:
39483946
vartype = type_map[expr]
39493947
return self.conditional_callable_type_map(expr, vartype)
@@ -3952,7 +3950,7 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39523950
# narrow their types. (For example, we shouldn't try narrowing the
39533951
# types of literal string or enum expressions).
39543952

3955-
operands = node.operands
3953+
operands = [collapse_walrus(x) for x in node.operands]
39563954
operand_types = []
39573955
narrowable_operand_index_to_hash = {}
39583956
for i, expr in enumerate(operands):
@@ -5742,3 +5740,14 @@ def has_bool_item(typ: ProperType) -> bool:
57425740
return any(is_named_instance(item, 'builtins.bool')
57435741
for item in typ.items)
57445742
return False
5743+
5744+
5745+
def collapse_walrus(e: Expression) -> Expression:
5746+
"""If an expression is an AssignmentExpr, pull out the assignment target.
5747+
5748+
We don't make any attempt to pull out all the targets in code like `x := (y := z)`.
5749+
We could support narrowing those if that sort of code turns out to be common.
5750+
"""
5751+
if isinstance(e, AssignmentExpr):
5752+
return e.target
5753+
return e

test-data/unit/check-python38.test

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ def f(p1: bytes, p2: float, /) -> None:
189189

190190
[case testWalrus]
191191
# flags: --strict-optional
192-
from typing import NamedTuple, Optional
192+
from typing import NamedTuple, Optional, List
193193
from typing_extensions import Final
194194

195195
if a := 2:
@@ -288,10 +288,23 @@ def check_partial() -> None:
288288

289289
reveal_type(x) # N: Revealed type is 'Union[builtins.int, None]'
290290

291-
def check_narrow(x: Optional[int]) -> None:
291+
def check_narrow(x: Optional[int], s: List[int]) -> None:
292292
if (y := x):
293293
reveal_type(y) # N: Revealed type is 'builtins.int'
294-
[builtins fixtures/f_string.pyi]
294+
295+
if (y := x) is not None:
296+
reveal_type(y) # N: Revealed type is 'builtins.int'
297+
298+
if (y := x) == 10:
299+
reveal_type(y) # N: Revealed type is 'builtins.int'
300+
301+
if (y := x) in s:
302+
reveal_type(y) # N: Revealed type is 'builtins.int'
303+
304+
if isinstance((y := x), int):
305+
reveal_type(y) # N: Revealed type is 'builtins.int'
306+
307+
[builtins fixtures/isinstancelist.pyi]
295308

296309
[case testWalrusPartialTypes]
297310
from typing import List

0 commit comments

Comments
 (0)