Skip to content

Commit d2b0234

Browse files
hauntsaninjahauntsaninja
authored
Narrow membership in statically known containers (#21461)
The negative narrowing here could be more aggressive, but we will need better literal handling to unblock it Builds on #21456 Fixes #13684 in combination with previous PRs. There is one remaining diagnostic, but that one is desirable
1 parent db331b4 commit d2b0234

3 files changed

Lines changed: 60 additions & 11 deletions

File tree

‎mypy/checker.py‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6772,14 +6772,34 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
67726772

67736773
if left_index in narrowable_operand_index_to_hash:
67746774
p_iterable_type = get_proper_type(iterable_type)
6775+
container_item_types = None
6776+
6777+
# Can we statically determine container contents?
67756778
if (
67766779
isinstance(p_iterable_type, TupleType)
67776780
and find_unpack_in_list(p_iterable_type.items) is None
67786781
):
6779-
# For some tuples, we can do negative narrowing, e.g. `x not in (None,)`
6782+
container_item_types = p_iterable_type.items
6783+
else:
6784+
container_expr = collapse_walrus(operands[right_index])
6785+
if isinstance(container_expr, (ListExpr, SetExpr)):
6786+
if all(not isinstance(i, StarExpr) for i in container_expr.items):
6787+
container_item_types = [
6788+
self.lookup_type(e) for e in container_expr.items
6789+
]
6790+
elif isinstance(container_expr, DictExpr):
6791+
if all(k is not None for k, v in container_expr.items):
6792+
container_item_types = [
6793+
self.lookup_type(cast(Expression, k))
6794+
for k, v in container_expr.items
6795+
]
6796+
6797+
if container_item_types is not None:
6798+
# If we know the exact contents, we can potentially do negative narrowing,
6799+
# e.g. `x not in (None,)`
67806800
all_if_maps = []
67816801
all_else_maps = []
6782-
for known_item in p_iterable_type.items:
6802+
for known_item in container_item_types:
67836803
# Match the should_coerce_literals logic from narrow_type_by_identity_equality
67846804
p_known_item = get_proper_type(known_item)
67856805
if is_literal_type_like(p_known_item) or (
@@ -6799,7 +6819,7 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
67996819
if_map = reduce_or_conditional_type_maps(all_if_maps)
68006820
else_map = reduce_and_conditional_type_maps(all_else_maps, use_meet=True)
68016821
else:
6802-
collection_item_type = get_proper_type(builtin_item_type(iterable_type))
6822+
collection_item_type = get_proper_type(builtin_item_type(p_iterable_type))
68036823
if collection_item_type is not None:
68046824
if_map, else_map = self.narrow_type_by_identity_equality(
68056825
"==",
@@ -6813,6 +6833,7 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
68136833
else_map = {}
68146834

68156835
if right_index in narrowable_operand_index_to_hash:
6836+
# E.g. narrows the right operand in `if "key" in typed_dict`
68166837
if_type, else_type = self.conditional_types_for_iterable(
68176838
item_type, iterable_type
68186839
)

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2254,7 +2254,6 @@ if y not in x:
22542254
else:
22552255
reveal_type(y) # N: Revealed type is "builtins.int"
22562256
[builtins fixtures/list.pyi]
2257-
[out]
22582257

22592258
[case testNarrowTypeAfterInListOfOptional]
22602259
# flags: --warn-unreachable
@@ -2268,7 +2267,6 @@ if y not in x:
22682267
else:
22692268
reveal_type(y) # N: Revealed type is "builtins.int | None"
22702269
[builtins fixtures/list.pyi]
2271-
[out]
22722270

22732271
[case testNarrowTypeAfterInListNonOverlapping]
22742272
# flags: --warn-unreachable
@@ -2319,7 +2317,6 @@ if y not in nt:
23192317
else:
23202318
reveal_type(y) # N: Revealed type is "builtins.int"
23212319
[builtins fixtures/tuple.pyi]
2322-
[out]
23232320

23242321
[case testNarrowTypeAfterInDict]
23252322
# flags: --warn-unreachable
@@ -2336,7 +2333,6 @@ if y not in x:
23362333
else:
23372334
reveal_type(y) # N: Revealed type is "builtins.str"
23382335
[builtins fixtures/dict.pyi]
2339-
[out]
23402336

23412337
[case testNarrowTypeAfterInNoAnyOrObject]
23422338
# flags: --warn-unreachable
@@ -2356,7 +2352,6 @@ else:
23562352
reveal_type(y) # N: Revealed type is "builtins.int | None"
23572353
[typing fixtures/typing-medium.pyi]
23582354
[builtins fixtures/list.pyi]
2359-
[out]
23602355

23612356
[case testNarrowTypeAfterInUserDefined]
23622357
# flags: --warn-unreachable
@@ -2378,7 +2373,6 @@ else:
23782373
reveal_type(y) # N: Revealed type is "builtins.int | None"
23792374
[typing fixtures/typing-full.pyi]
23802375
[builtins fixtures/list.pyi]
2381-
[out]
23822376

23832377
[case testNarrowTypeAfterInSet]
23842378
# flags: --warn-unreachable
@@ -2395,7 +2389,6 @@ if y not in s:
23952389
else:
23962390
reveal_type(y) # N: Revealed type is "builtins.str"
23972391
[builtins fixtures/set.pyi]
2398-
[out]
23992392

24002393
[case testNarrowTypeAfterInTypedDict]
24012394
# flags: --warn-unreachable
@@ -2412,7 +2405,6 @@ def f() -> None:
24122405
reveal_type(x) # N: Revealed type is "builtins.str"
24132406
[builtins fixtures/dict.pyi]
24142407
[typing fixtures/typing-typeddict.pyi]
2415-
[out]
24162408

24172409
[case testIsinstanceWidensWithAnyArg]
24182410
# flags: --warn-unreachable

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3291,11 +3291,47 @@ def narrow_list(x: Literal['a', 'b', 'c'], t: list[Literal['a', 'b']]):
32913291
else:
32923292
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
32933293

3294+
if x in ['a', 'b']:
3295+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b']"
3296+
else:
3297+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3298+
3299+
if x in ['a', 'b', *[]]:
3300+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3301+
else:
3302+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3303+
32943304
def narrow_set(x: Literal['a', 'b', 'c'], t: set[Literal['a', 'b']]):
32953305
if x in t:
32963306
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b']"
32973307
else:
32983308
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3309+
3310+
if x in {'a', 'b'}:
3311+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b']"
3312+
else:
3313+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3314+
3315+
if x in {'a', 'b', *[]}:
3316+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3317+
else:
3318+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3319+
3320+
def narrow_dict(x: Literal['a', 'b', 'c'], t: dict[Literal['a', 'b'], int]):
3321+
if x in t: # E: Unsupported operand types for in ("Literal['a', 'b', 'c']" and "dict[Literal['a', 'b'], int]")
3322+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b']"
3323+
else:
3324+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3325+
3326+
if x in {'a': 0, 'b': 1}:
3327+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b']"
3328+
else:
3329+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3330+
3331+
if x in {'a': 0, 'b': 1, **{}}:
3332+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
3333+
else:
3334+
reveal_type(x) # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
32993335
[builtins fixtures/primitives.pyi]
33003336

33013337

0 commit comments

Comments
 (0)