Skip to content

Commit c4d8554

Browse files
committed
Support narrowing of walrus in most cases (#8458)
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 dc0d35f commit c4d8554

File tree

2 files changed

+31
-7
lines changed

2 files changed

+31
-7
lines changed

mypy/checker.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3924,10 +3924,12 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39243924
elif is_false_literal(node):
39253925
return None, {}
39263926
elif isinstance(node, CallExpr):
3927+
if len(node.args) == 0:
3928+
return {}, {}
3929+
expr = collapse_walrus(node.args[0])
39273930
if refers_to_fullname(node.callee, 'builtins.isinstance'):
39283931
if len(node.args) != 2: # the error will be reported elsewhere
39293932
return {}, {}
3930-
expr = node.args[0]
39313933
if literal(expr) == LITERAL_TYPE:
39323934
return self.conditional_type_map_with_intersection(
39333935
expr,
@@ -3937,13 +3939,11 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39373939
elif refers_to_fullname(node.callee, 'builtins.issubclass'):
39383940
if len(node.args) != 2: # the error will be reported elsewhere
39393941
return {}, {}
3940-
expr = node.args[0]
39413942
if literal(expr) == LITERAL_TYPE:
39423943
return self.infer_issubclass_maps(node, expr, type_map)
39433944
elif refers_to_fullname(node.callee, 'builtins.callable'):
39443945
if len(node.args) != 1: # the error will be reported elsewhere
39453946
return {}, {}
3946-
expr = node.args[0]
39473947
if literal(expr) == LITERAL_TYPE:
39483948
vartype = type_map[expr]
39493949
return self.conditional_callable_type_map(expr, vartype)
@@ -3952,7 +3952,7 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
39523952
# narrow their types. (For example, we shouldn't try narrowing the
39533953
# types of literal string or enum expressions).
39543954

3955-
operands = node.operands
3955+
operands = [collapse_walrus(x) for x in node.operands]
39563956
operand_types = []
39573957
narrowable_operand_index_to_hash = {}
39583958
for i, expr in enumerate(operands):
@@ -5742,3 +5742,14 @@ def has_bool_item(typ: ProperType) -> bool:
57425742
return any(is_named_instance(item, 'builtins.bool')
57435743
for item in typ.items)
57445744
return False
5745+
5746+
5747+
def collapse_walrus(e: Expression) -> Expression:
5748+
"""If an expression is an AssignmentExpr, pull out the assignment target.
5749+
5750+
We don't make any attempt to pull out all the targets in code like `x := (y := z)`.
5751+
We could support narrowing those if that sort of code turns out to be common.
5752+
"""
5753+
if isinstance(e, AssignmentExpr):
5754+
return e.target
5755+
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)