Improve negative narrowing for membership checks on tuples (#21456)

Related to #21411

Fixes #16093

Author: hauntsaninja

mypy/checker.py

@@ -6771,25 +6771,45 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
                 else_map = {}
 
                 if left_index in narrowable_operand_index_to_hash:
-                    collection_item_type = get_proper_type(builtin_item_type(iterable_type))
-                    if collection_item_type is not None:
-                        if_map, else_map = self.narrow_type_by_identity_equality(
-                            "==",
-                            operands=[operands[left_index], operands[right_index]],
-                            operand_types=[item_type, collection_item_type],
-                            expr_indices=[0, 1],
-                            narrowable_indices={0},
-                        )
-                        if else_map and not (
-                            isinstance(p_typ := get_proper_type(iterable_type), TupleType)
-                            and all(
-                                is_singleton_equality_type(get_proper_type(item))
-                                for item in p_typ.items
+                    p_iterable_type = get_proper_type(iterable_type)
+                    if (
+                        isinstance(p_iterable_type, TupleType)
+                        and find_unpack_in_list(p_iterable_type.items) is None
+                    ):
+                        # For some tuples, we can do negative narrowing, e.g. `x not in (None,)`
+                        all_if_maps = []
+                        all_else_maps = []
+                        for known_item in p_iterable_type.items:
+                            # Match the should_coerce_literals logic from narrow_type_by_identity_equality
+                            p_known_item = get_proper_type(known_item)
+                            if is_literal_type_like(p_known_item) or (
+                                isinstance(p_known_item, Instance) and p_known_item.type.is_enum
+                            ):
+                                known_item = coerce_to_literal(known_item)
+                            if_map, else_map = self.narrow_type_by_identity_equality(
+                                "==",
+                                operands=[operands[left_index], operands[right_index]],
+                                operand_types=[item_type, known_item],
+                                expr_indices=[0, 1],
+                                narrowable_indices={0},
                             )
-                        ):
-                            # In general, we can't do negative narrowing, since e.g. the container
-                            # could just be empty. However, we can do negative narrowing for some
-                            # tuples e.g. `x not in (None,)`
+                            all_if_maps.append(if_map)
+                            if is_singleton_equality_type(get_proper_type(known_item)):
+                                all_else_maps.append(else_map)
+                        if_map = reduce_or_conditional_type_maps(all_if_maps)
+                        else_map = reduce_and_conditional_type_maps(all_else_maps, use_meet=True)
+                    else:
+                        collection_item_type = get_proper_type(builtin_item_type(iterable_type))
+                        if collection_item_type is not None:
+                            if_map, else_map = self.narrow_type_by_identity_equality(
+                                "==",
+                                operands=[operands[left_index], operands[right_index]],
+                                operand_types=[item_type, collection_item_type],
+                                expr_indices=[0, 1],
+                                narrowable_indices={0},
+                            )
+                            # We can't do negative narrowing, since e.g. the container could
+                            # just be empty.
                             else_map = {}
 
                 if right_index in narrowable_operand_index_to_hash:

test-data/unit/check-isinstance.test

@@ -2294,18 +2294,16 @@ def f(x: Optional[int], lst: Optional[List[int]], nested_any: List[List[Any]]) -
 
 [case testNarrowTypeAfterInTuple]
 # flags: --warn-unreachable
-from typing import Optional
 class A: pass
 class B(A): pass
 class C(A): pass
 
-y: Optional[B]
-if y in (B(), C()):
-    reveal_type(y) # N: Revealed type is "__main__.B"
-else:
-    reveal_type(y) # N: Revealed type is "__main__.B | None"
+def f(y: B | None):
+    if y in (B(), C()):
+        reveal_type(y) # N: Revealed type is "__main__.B"
+    else:
+        reveal_type(y) # N: Revealed type is "__main__.B | None"
 [builtins fixtures/tuple.pyi]
-[out]
 
 [case testNarrowTypeAfterInNamedTuple]
 # flags: --warn-unreachable

test-data/unit/check-narrowing.test

@@ -3201,7 +3201,7 @@ class X:
 [builtins fixtures/dict.pyi]
 
 
-[case testTypeNarrowingStringInLiteralContainer]
+[case testNarrowStringInLiteralContainer]
 # flags: --strict-equality --warn-unreachable
 from typing import Literal
 
@@ -3235,6 +3235,69 @@ def narrow_set(x: str, t: set[Literal['a', 'b']]):
         reveal_type(x)  # N: Revealed type is "builtins.str"
 [builtins fixtures/primitives.pyi]
 
+[case testNarrowLiteralInLiteralContainer]
+# flags: --strict-equality --warn-unreachable
+from typing import Literal
+
+def narrow_tuple_exact(x: Literal['a', 'b', 'c'], t: tuple[Literal['a'], Literal['b']]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['c']"
+
+    if x not in t:
+        reveal_type(x)  # N: Revealed type is "Literal['c']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+
+def narrow_tuple_expression(x: Literal['a', 'b', 'c']):
+    # TODO: this should match narrow_tuple_exact
+    if x in ('a', 'b'):
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+
+    if x not in ('a', 'b'):
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+
+def narrow_tuple_union(x: Literal['a', 'b', 'c'], t: tuple[Literal['a', 'b']]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+
+    if x not in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+
+def narrow_tuple_with_other_type(x: Literal['a', 'b', 'c'], t: tuple[Literal['a'], int]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['b'] | Literal['c']"
+
+def narrow_homo_tuple(x: Literal['a', 'b', 'c'], t: tuple[Literal['a', 'b'], ...]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+
+def narrow_list(x: Literal['a', 'b', 'c'], t: list[Literal['a', 'b']]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+
+def narrow_set(x: Literal['a', 'b', 'c'], t: set[Literal['a', 'b']]):
+    if x in t:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b']"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal['a'] | Literal['b'] | Literal['c']"
+[builtins fixtures/primitives.pyi]
+
 
 [case testNarrowingLiteralInLiteralContainer]
 # flags: --strict-equality --warn-unreachable

test-data/unit/check-typevar-tuple.test

@@ -2145,14 +2145,13 @@ match(b)  # E: Argument 1 to "match" has incompatible type "Bad"; expected "PC[U
 [builtins fixtures/tuple.pyi]
 
 [case testVariadicTupleCollectionCheck]
-from typing import Tuple, Optional
 from typing_extensions import Unpack
 
-allowed: Tuple[int, Unpack[Tuple[int, ...]]]
+allowed: tuple[int, Unpack[tuple[int, ...]]]
 
-x: Optional[int]
-if x in allowed:
-    reveal_type(x)  # N: Revealed type is "builtins.int"
+def f(x: int | None):
+    if x in allowed:
+        reveal_type(x)  # N: Revealed type is "builtins.int"
 [builtins fixtures/tuple.pyi]
 
 [case testJoinOfVariadicTupleCallablesNoCrash]