Skip to content

Commit ce186f4

Browse files
authored
Allow 'in' operations with partial types (#8177)
This improves type inference in code like this: ``` d = {} for k in f(): if k not in d: d[k] = 0 else: d[k] += 1 ``` Unfortunately, this change breaks strict equality checking in 'in' operations involving partial types. I'm not 100% sure this is a net improvement. I couldn't find a clean way of getting strict equality working together with this.
1 parent 387a911 commit ce186f4

2 files changed

Lines changed: 49 additions & 1 deletion

File tree

‎mypy/checkexpr.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2125,7 +2125,11 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
21252125
method_type = None # type: Optional[mypy.types.Type]
21262126

21272127
if operator == 'in' or operator == 'not in':
2128-
right_type = self.accept(right) # always validate the right operand
2128+
# If the right operand has partial type, look it up without triggering
2129+
# a "Need type annotation ..." message, as it would be noise.
2130+
right_type = self.find_partial_type_ref_fast_path(right)
2131+
if right_type is None:
2132+
right_type = self.accept(right) # Validate the right operand
21292133

21302134
# Keep track of whether we get type check errors (these won't be reported, they
21312135
# are just to verify whether something is valid typing wise).
@@ -2206,6 +2210,22 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
22062210
assert result is not None
22072211
return result
22082212

2213+
def find_partial_type_ref_fast_path(self, expr: Expression) -> Optional[Type]:
2214+
"""If expression has a partial generic type, return it without additional checks.
2215+
2216+
In particular, this does not generate an error about a missing annotation.
2217+
2218+
Otherwise, return None.
2219+
"""
2220+
if not isinstance(expr, RefExpr):
2221+
return None
2222+
if isinstance(expr.node, Var):
2223+
result = self.analyze_var_ref(expr.node, expr)
2224+
if isinstance(result, PartialType) and result.type is not None:
2225+
self.chk.store_type(expr, self.chk.fixup_partial_type(result))
2226+
return result
2227+
return None
2228+
22092229
def dangerous_comparison(self, left: Type, right: Type,
22102230
original_container: Optional[Type] = None) -> bool:
22112231
"""Check for dangerous non-overlapping comparisons like 42 == 'no'.

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1738,6 +1738,34 @@ def g() -> None:
17381738
reveal_type(y) # N: Revealed type is 'builtins.list[Any]'
17391739
[builtins fixtures/list.pyi]
17401740

1741+
[case testInferFromEmptyDictWhenUsingIn]
1742+
d = {}
1743+
if 'x' in d:
1744+
d['x'] = 1
1745+
reveal_type(d) # N: Revealed type is 'builtins.dict[builtins.str, builtins.int]'
1746+
1747+
dd = {}
1748+
if 'x' not in dd:
1749+
dd['x'] = 1
1750+
reveal_type(dd) # N: Revealed type is 'builtins.dict[builtins.str, builtins.int]'
1751+
[builtins fixtures/dict.pyi]
1752+
1753+
[case testInferFromEmptyDictWhenUsingInSpecialCase]
1754+
d = None
1755+
if 'x' in d: # E: "None" has no attribute "__iter__" (not iterable)
1756+
pass
1757+
reveal_type(d) # N: Revealed type is 'None'
1758+
[builtins fixtures/dict.pyi]
1759+
1760+
[case testInferFromEmptyListWhenUsingInWithStrictEquality]
1761+
# flags: --strict-equality
1762+
def f() -> None:
1763+
a = []
1764+
if 1 in a: # TODO: This should be an error
1765+
a.append('x')
1766+
[builtins fixtures/list.pyi]
1767+
[typing fixtures/typing-full.pyi]
1768+
17411769

17421770
-- Inferring types of variables first initialized to None (partial types)
17431771
-- ----------------------------------------------------------------------

0 commit comments

Comments
 (0)