8000 Use type information from isinstance checks in comprehensions (#2000) · python/mypy@e2561a8 · GitHub
[go: up one dir, main page]
Skip to content
Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit e2561a8

Browse files
ddfishergvanrossum
ddfisher
authored and
gvanrossum
committed
Use type information from isinstance checks in comprehensions (#2000)
Fix #1734.
1 parent d83aece commit e2561a8

File tree

4 files changed

+77
-43
lines changed

4 files changed

+77
-43
lines changed

‎mypy/checkexpr.py

Lines changed: 53 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1531,55 +1531,68 @@ def check_generator_or_comprehension(self, gen: GeneratorExpr,
15311531
type_name: str,
15321532
id_for_messages: str) -> Type:
15331533
"""Type check a generator expression or a list comprehension."""
1534-
self.check_for_comp(gen)
1534+
with self.chk.binder.frame_context():
1535+
self.check_for_comp(gen)
15351536

1536-
# Infer the type of the list comprehension by using a synthetic generic
1537-
# callable type.
1538-
tvdef = TypeVarDef('T', -1, [], self.chk.object_type())
1539-
tv = TypeVarType(tvdef)
1540-
constructor = CallableType(
1541-
[tv],
1542-
[nodes.ARG_POS],
1543-
[None],
1544-
self.chk.named_generic_type(type_name, [tv]),
1545-
self.chk.named_type('builtins.function'),
1546-
name=id_for_messages,
1547-
variables=[tvdef])
1548-
return self.check_call(constructor,
1549-
[gen.left_expr], [nodes.ARG_POS], gen)[0]
1537+
# Infer the type of the list comprehension by using a synthetic generic
1538+
# callable type.
1539+
tvdef = TypeVarDef('T', -1, [], self.chk.object_type())
1540+
tv = TypeVarType(tvdef)
1541+
constructor = CallableType(
1542+
[tv],
1543+
[nodes.ARG_POS],
1544+
[None],
1545+
self.chk.named_generic_type(type_name, [tv]),
1546+
self.chk.named_type('builtins.function'),
1547+
name=id_for_messages,
1548+
variables=[tvdef])
1549+
return self.check_call(constructor,
1550+
[gen.left_expr], [nodes.ARG_POS], gen)[0]
15501551

15511552
def visit_dictionary_comprehension(self, e: DictionaryComprehension) -> Type:
15521553
"""Type check a dictionary comprehension."""
1553-
self.check_for_comp(e)
1554-
1555-
# Infer the type of the list comprehension by using a synthetic generic
1556-
# callable type.
1557-
ktdef = TypeVarDef('KT', -1, [], self.chk.object_type())
1558-
vtdef = TypeVarDef('VT', -2, [], self.chk.object_type())
1559-
kt = TypeVarType(ktdef)
1560-
vt = TypeVarType(vtdef)
1561-
constructor = CallableType(
1562-
[kt, vt],
1563-
[nodes.ARG_POS, nodes.ARG_POS],
1564-
[None, None],
1565-
self.chk.named_generic_type('builtins.dict', [kt, vt]),
1566-
self.chk.named_type('builtins.function'),
1567-
name='<dictionary-comprehension>',
1568-
variables=[ktdef, vtdef])
1569-
return self.check_call(constructor,
1570-
[e.key, e.value], [nodes.ARG_POS, nodes.ARG_POS], e)[0]
1554+
with self.chk.binder.frame_context():
1555+
self.check_for_comp(e)
1556+
1557+
# Infer the type of the list comprehension by using a synthetic generic
1558+
# callable type.
1559+
ktdef = TypeVarDef('KT', -1, [], self.chk.object_type())
1560+
vtdef = TypeVarDef('VT', -2, [], self.chk.object_type())
1561+
kt = TypeVarType(ktdef)
1562+
vt = TypeVarType(vtdef)
1563+
constructor = CallableType(
1564+
[kt, vt],
1565+
[nodes.ARG_POS, nodes.ARG_POS],
1566+
[None, None],
1567+
self.chk.named_generic_type('builtins.dict', [kt, vt]),
1568+
self.chk.named_type('builtins.function'),
1569+
name='<dictionary-comprehension>',
1570+
variables=[ktdef, vtdef])
1571+
return self.check_call(constructor,
1572+
[e.key, e.value], [nodes.ARG_POS, nodes.ARG_POS], e)[0]
15711573

15721574
def check_for_comp(self, e: Union[GeneratorExpr, DictionaryComprehension]) -> None:
15731575
"""Check the for_comp part of comprehensions. That is the part from 'for':
15741576
... for x in y if z
1577+
1578+
Note: This adds the type information derived from the condlists to the current binder.
15751579
"""
1576-
with self.chk.binder.frame_context():
1577-
for index, sequence, conditions in zip(e.indices, e.sequences,
1578-
e.condlists):
1579-
sequence_type = self.chk.analyze_iterable_item_type(sequence)
1580-
self.chk.analyze_index_variables(index, sequence_type, e)
1581-
for condition in conditions:
1582-
self.accept(condition)
1580+
for index, sequence, conditions in zip(e.indices, e.sequences,
1581+
e.condlists):
1582+
sequence_type = self.chk.analyze_iterable_item_type(sequence)
1583+
self.chk.analyze_index_variables(index, sequence_type, e)
1584+
for condition in conditions:
1585+
self.< 1E0A span class=pl-c1>accept(condition)
1586+
1587+
# values are only part of the comprehension when all conditions are true
1588+
true_map, _ = mypy.checker.find_isinstance_check(
1589+
condition, self.chk.type_map,
1590+
self.chk.typing_mode_weak()
1591+
)
1592+
1593+
if true_map:
1594+
for var, type in true_map.items():
1595+
self.chk.binder.push(var, type)
15831596

15841597
def visit_conditional_expr(self, e: ConditionalExpr) -> Type:
15851598
cond_type = self.accept(e.cond)

‎mypy/nodes.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1504,7 +1504,7 @@ class GeneratorExpr(Expression):
15041504
"""Generator expression ... for ... in ... [ for ... in ... ] [ if ... ]."""
15051505

15061506
left_expr = None # type: Expression
1507-
sequences_expr = None # type: List[Expression]
1507+
sequences = None # type: List[Expression]
15081508
condlists = None # type: List[List[Expression]]
15091509
indices = None # type: List[Expression]
15101510

@@ -1548,7 +1548,7 @@ class DictionaryComprehension(Expression):
15481548

15491549
key = None # type: Expression
15501550
value = None # type: Expression
1551-
sequences_expr = None # type: List[Expression]
1551+
sequences = None # type: List[Expression]
15521552
condlists = None # type: List[List[Expression]]
15531553
indices = None # type: List[Expression]
15541554

‎test-data/unit/check-isinstance.test

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1155,3 +1155,13 @@ else:
11551155
1()
11561156
[builtins fixtures/isinstance.py]
11571157
[out]
1158+
[case testComprehensionIsInstance]
1159+
from typing import List, Union
1160+
a = [] # type: List[Union[int, str]]
1161+
l = [x for x in a if isinstance(x, int)]
1162+
g = (x for x in a if isinstance(x, int))
1163+
d = {0: x for x in a if isinstance(x, int)}
1164+
reveal_type(l) # E: Revealed type is 'builtins.list[builtins.int*]'
1165+
reveal_type(g) # E: Revealed type is 'typing.Iterator[builtins.int*]'
1166+
reveal_type(d) # E: Revealed type is 'builtins.dict[builtins.int*, builtins.int*]'
1167+
[builtins fixtures/isinstancelist.py]

‎test-data/unit/fixtures/isinstancelist.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import builtinclass, Iterable, Iterator, Generic, TypeVar, List
1+
from typing import builtinclass, Iterable, Iterator, Generic, TypeVar, List, Mapping, overload, Tuple
22

33
@builtinclass
44
class object:
@@ -24,10 +24,21 @@ def __add__(self, x: str) -> str: pass
2424
def __getitem__(self, x: int) -> str: pass
2525

2626
T = TypeVar('T')
27+
KT = TypeVar('KT')
28+
VT = TypeVar('VT')
2729

2830
class list(Iterable[T], Generic[T]):
2931
def __iter__(self) -> Iterator[T]: pass
3032
def __mul__(self, x: int) -> list[T]: pass
3133
def __setitem__(self, x: int, v: T) -> None: pass
3234
def __getitem__(self, x: int) -> T: pass
3335
def __add__(self, x: List[T]) -> T: pass
36+
37+
class dict(Iterable[KT], Mapping[KT, VT], Generic[KT, VT]):
38+
@overload
39+
def __init__(self, **kwargs: VT) -> None: pass
40+
@overload
41+
def __init__(self, arg: Iterable[Tuple[KT, VT]], **kwargs: VT) -> None: pass
42+
def __setitem__(self, k: KT, v: VT) -> None: pass
43+
def __iter__(self) -> Iterator[KT]: pass
44+
def update(self, a: Mapping[KT, VT]) -> None: pass

0 commit comments

Comments
 (0)
0