@@ -158,8 +158,10 @@ def try_infer_partial_type(self, e: CallExpr) -> None:
158
158
var = cast (Var , e .callee .expr .node )
159
159
partial_types = self .chk .find_partial_types (var )
160
160
if partial_types is not None and not self .chk .current_node_deferred :
161
- partial_type = cast (PartialType , var .type )
162
- if partial_type is None or partial_type .type is None :
161
+ partial_type = var .type
162
+ if (partial_type is None or
163
+ not isinstance (partial_type , PartialType ) or
164
+ partial_type .type is None ):
163
165
# A partial None type -> can't infer anything.
164
166
return
165
167
typename = partial_type .type .fullname ()
@@ -1538,55 +1540,68 @@ def check_generator_or_comprehension(self, gen: GeneratorExpr,
1538
1540
type_name : str ,
1539
1541
id_for_messages : str ) -> Type :
1540
1542
"""Type check a generator expression or a list comprehension."""
1541
- self .check_for_comp (gen )
1543
+ with self .chk .binder .frame_context ():
1544
+ self .check_for_comp (gen )
1542
1545
1543
- # Infer the type of the list comprehension by using a synthetic generic
1544
- # callable type.
1545
- tvdef = TypeVarDef ('T' , - 1 , [], self .chk .object_type ())
1546
- tv = TypeVarType (tvdef )
1547
- constructor = CallableType (
1548
- [tv ],
1549
- [nodes .ARG_POS ],
1550
- [None ],
1551
- self .chk .named_generic_type (type_name , [tv ]),
1552
- self .chk .named_type ('builtins.function' ),
1553
- name = id_for_messages ,
1554
- variables = [tvdef ])
1555
- return self .check_call (constructor ,
1556
- [gen .left_expr ], [nodes .ARG_POS ], gen )[0 ]
1546
+ # Infer the type of the list comprehension by using a synthetic generic
1547
+ # callable type.
1548
+ tvdef = TypeVarDef ('T' , - 1 , [], self .chk .object_type ())
1549
+ tv = TypeVarType (tvdef )
1550
+ constructor = CallableType (
1551
+ [tv ],
1552
+ [nodes .ARG_POS ],
1553
+ [None ],
1554
+ self .chk .named_generic_type (type_name , [tv ]),
1555
+ self .chk .named_type ('builtins.function' ),
1556
+ name = id_for_messages ,
1557
+ variables = [tvdef ])
1558
+ return self .check_call (constructor ,
1559
+ [gen .left_expr ], [nodes .ARG_POS ], gen )[0 ]
1557
1560
1558
1561
def visit_dictionary_comprehension (self , e : DictionaryComprehension ) -> Type :
1559
1562
"""Type check a dictionary comprehension."""
1560
- self .check_for_comp (e )
1561
-
1562
- # Infer the type of the list comprehension by using a synthetic generic
1563
- # callable type.
1564
- ktdef = TypeVarDef ('KT' , - 1 , [], self .chk .object_type ())
1565
- vtdef = TypeVarDef ('VT' , - 2 , [], self .chk .object_type ())
1566
- kt = TypeVarType (ktdef )
1567
- vt = TypeVarType (vtdef )
1568
- constructor = CallableType (
1569
- [kt , vt ],
1570
- [nodes .ARG_POS , nodes .ARG_POS ],
1571
- [None , None ],
1572
- self .chk .named_generic_type ('builtins.dict' , [kt , vt ]),
1573
- self .chk .named_type ('builtins.function' ),
1574
- name = '<dictionary-comprehension>' ,
1575
- variables = [ktdef , vtdef ])
1576
- return self .check_call (constructor ,
1577
- [e .key , e .value ], [nodes .ARG_POS , nodes .ARG_POS ], e )[0 ]
1563
+ with self .chk .binder .frame_context ():
1564
+ self .check_for_comp (e )
1565
+
1566
+ # Infer the type of the list comprehension by using a synthetic generic
1567
+ # callable type.
1568
+ ktdef = TypeVarDef ('KT' , - 1 , [], self .chk .object_type ())
1569
+ vtdef = TypeVarDef ('VT' , - 2 , [], self .chk .object_type ())
1570
+ kt = TypeVarType (ktdef )
1571
+ vt = TypeVarType (vtdef )
1572
+ constructor = CallableType (
1573
+ [kt , vt ],
1574
+ [nodes .ARG_POS , nodes .ARG_POS ],
1575
+ [None , None ],
1576
+ self .chk .named_generic_type ('builtins.dict' , [kt , vt ]),
1577
+ self .chk .named_type ('builtins.function' ),
1578
+ name = '<dictionary-comprehension>' ,
1579
+ variables = [ktdef , vtdef ])
1580
+ return self .check_call (constructor ,
1581
+ [e .key , e .value ], [nodes .ARG_POS , nodes .ARG_POS ], e )[0 ]
1578
1582
1579
1583
def check_for_comp (self , e : Union [GeneratorExpr , DictionaryComprehension ]) -> None :
1580
1584
"""Check the for_comp part of comprehensions. That is the part from 'for':
1581
1585
... for x in y if z
1586
+
1587
+ Note: This adds the type information derived from the condlists to the current binder.
1582
1588
"""
1583
- with self .chk .binder .frame_context ():
1584
- for index , sequence , conditions in zip (e .indices , e .sequences ,
1585
- e .condlists ):
1586
- sequence_type = self .chk .analyze_iterable_item_type (sequence )
1587
- self .chk .analyze_index_variables (index , sequence_type , e )
1588
- for condition in conditions :
1589
- self .accept (condition )
1589
+ for index , sequence , conditions in zip (e .indices , e .sequences ,
1590
+ e .condlists ):
1591
+ sequence_type = self .chk .analyze_iterable_item_type (sequence )
1592
+ self .chk .analyze_index_variables (index , sequence_type , e )
1593
+ for condition in conditions :
1594
+ self .accept (condition )
1595
+
1596
+ # values are only part of the comprehension when all conditions are true
1597
+ true_map , _ = mypy .checker .find_isinstance_check (
1598
+ condition , self .chk .type_map ,
1599
+ self .chk .typing_mode_weak ()
1600
+ )
1601
+
1602
+ if true_map :
1603
+ for var , type in true_map .items ():
1604
+ self .chk .binder .push (var , type )
1590
1605
1591
1606
def visit_conditional_expr (self , e : ConditionalExpr ) -> Type :
1592
1607
cond_type = self .accept (e .cond )
0 commit comments