Support exhaustive checking for “in” operator.

feature-test.py

@@ -0,0 +1,19 @@
+from enum import Enum
+from typing_extensions import assert_never
+
+
+class MyEnum(Enum):
+    A = 1
+    B = 2
+    C = 3
+
+
+def my_function(a: MyEnum) -> bool:
+    if a == MyEnum.A:
+        return True
+    elif a in (MyEnum.B, MyEnum.C):
+        return False
+    assert_never(a)
+
+
+my_function(MyEnum.A)

mypy/checker.py

@@ -1247,7 +1247,8 @@ def check_func_def(
                         if (
                             arg_type.variance == COVARIANT
                             and defn.name not in ("__init__", "__new__", "__post_init__")
-                            and not is_private(defn.name)  # private methods are not inherited
+                            # private methods are not inherited
+                            and not is_private(defn.name)
                         ):
                             ctx: Context = arg_type
                             if ctx.line < 0:
@@ -3112,7 +3113,8 @@ def check_compatibility_all_supers(
         if (
             isinstance(lvalue_node, Var)
             and lvalue.kind in (MDEF, None)
-            and len(lvalue_node.info.bases) > 0  # None for Vars defined via self
+            # None for Vars defined via self
+            and len(lvalue_node.info.bases) > 0
         ):
             for base in lvalue_node.info.mro[1:]:
                 tnode = base.names.get(lvalue_node.name)
@@ -5708,6 +5710,7 @@ def find_isinstance_check_helper(self, node: Expression) -> tuple[TypeMap, TypeM
             operands = [collapse_walrus(x) for x in node.operands]
             operand_types = []
             narrowable_operand_index_to_hash = {}
+            # print(operands)
             for i, expr in enumerate(operands):
                 if not self.has_type(expr):
                     return {}, {}
@@ -5810,49 +5813,62 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                             narrowable_operand_index_to_hash.keys(),
                         )
 
-                    # If we haven't been able to narrow types yet, we might be dealing with a
-                    # explicit type(x) == some_type check
+                    # If we haven't been able to narrow types yet, we might be dealing with a explicit type(x) == some_type check
                     if if_map == {} and else_map == {}:
                         if_map, else_map = self.find_type_equals_check(node, expr_indices)
+                        # print("done4", if_map, else_map)
                 elif operator in {"in", "not in"}:
                     assert len(expr_indices) == 2
                     left_index, right_index = expr_indices
                     item_type = operand_types[left_index]
                     iterable_type = operand_types[right_index]
+                    right_expr = operands[right_index]
 
                     if_map, else_map = {}, {}
 
-                    if left_index in narrowable_operand_index_to_hash:
-                        # We only try and narrow away 'None' for now
-                        if is_overlapping_none(item_type):
-                            collection_item_type = get_proper_type(
-                                builtin_item_type(iterable_type)
-                            )
-                            if (
-                                collection_item_type is not None
-                                and not is_overlapping_none(collection_item_type)
-                                and not (
-                                    isinstance(collection_item_type, Instance)
-                                    and collection_item_type.type.fullname == "builtins.object"
-                                )
-                                and is_overlapping_erased_types(item_type, collection_item_type)
-                            ):
-                                if_map[operands[left_index]] = remove_optional(item_type)
-
-                    if right_index in narrowable_operand_index_to_hash:
-                        if_type, else_type = self.conditional_types_for_iterable(
-                            item_type, iterable_type
+                    if isinstance(right_expr, TupleExpr):
+                        all_literal_enum = all(
+                            self.is_literal_enum(element_expr) for element_expr in right_expr.items
                         )
-                        expr = operands[right_index]
-                        if if_type is None:
-                            if_map = None
-                        else:
-                            if_map[expr] = if_type
-                        if else_type is None:
+                        if all_literal_enum:
+                            # Set if_map for the entire tuple
+                            if_map = {}
                             else_map = None
-                        else:
-                            else_map[expr] = else_type
 
+                    else:
+                        if left_index in narrowable_operand_index_to_hash:
+                            # We only try and narrow away 'None' for now
+                            if is_overlapping_none(item_type):
+                                collection_item_type = get_proper_type(
+                                    builtin_item_type(iterable_type)
+                                )
+                                if (
+                                    collection_item_type is not None
+                                    and not is_overlapping_none(collection_item_type)
+                                    and not (
+                                        isinstance(collection_item_type, Instance)
+                                        and collection_item_type.type.fullname == "builtins.object"
+                                    )
+                                    and is_overlapping_erased_types(
+                                        item_type, collection_item_type
+                                    )
+                                ):
+                                    if_map[operands[left_index]] = remove_optional(item_type)
+
+                        if right_index in narrowable_operand_index_to_hash:
+                            if_type, else_type = self.conditional_types_for_iterable(
+                                item_type, iterable_type
+                            )
+                            expr = operands[right_index]
+
+                            if if_type is None:
+                                if_map = None
+                            else:
+                                if_map[expr] = if_type
+                            if else_type is None:
+                                else_map = None
+                            else:
+                                else_map[expr] = else_type
                 else:
                     if_map = {}
                     else_map = {}
@@ -6151,6 +6167,7 @@ def refine_identity_comparison_expression(
         expressions in the chain to a Literal type. Performing this coercion is sometimes
         too aggressive of a narrowing, depending on context.
         """
+
         should_coerce = True
         if coerce_only_in_literal_context:
 
@@ -6246,9 +6263,6 @@ def should_coerce_inner(typ: Type) -> bool:
             if sum_type_name is not None:
                 expr_type = try_expanding_sum_type_to_union(expr_type, sum_type_name)
 
-            # We intentionally use 'conditional_types' directly here instead of
-            # 'self.conditional_types_with_intersection': we only compute ad-hoc
-            # intersections when working with pure instances.
             types = conditional_types(expr_type, target_type)
             partial_type_maps.append(conditional_types_to_typemaps(expr, *types))
 
@@ -7196,7 +7210,10 @@ class Foo(Enum):
         unit for the same reasons we sometimes treat 'True', 'False', or 'None' as a single
         primitive unit.
         """
+
         if not isinstance(n, MemberExpr) or not isinstance(n.expr, NameExpr):
+            # print(n, isinstance(n, MemberExpr))
+            # print(n, isinstance(n.expr, NameExpr))
             return False
 
         parent_type = self.lookup_type_or_none(n.expr)
@@ -7502,7 +7519,8 @@ def builtin_item_type(tp: Type) -> Type | None:
             else:
                 normalized_items.append(it)
         if all(not isinstance(it, AnyType) for it in get_proper_types(normalized_items)):
-            return make_simplified_union(normalized_items)  # this type is not externally visible
+            # this type is not externally visible
+            return make_simplified_union(normalized_items)
     elif isinstance(tp, TypedDictType):
         # TypedDict always has non-optional string keys. Find the key type from the Mapping
         # base class.

mypy/checkexpr.py

@@ -3473,7 +3473,8 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
                 # a "Need type annotation ..." message, as it would be noise.
                 right_type = self.find_partial_type_ref_fast_path(right)
                 if right_type is None:
-                    right_type = self.accept(right)  # Validate the right operand
+                    # Validate the right operand
+                    right_type = self.accept(right)
 
                 right_type = get_proper_type(right_type)
                 item_types: Sequence[Type] = [right_type]
@@ -6094,7 +6095,8 @@ def __init__(self, ignore_in_type_obj: bool) -> None:
         self.ignore_in_type_obj = ignore_in_type_obj
 
     def visit_any(self, t: AnyType) -> bool:
-        return t.type_of_any != TypeOfAny.special_form  # special forms are not real Any types
+        # special forms are not real Any types
+        return t.type_of_any != TypeOfAny.special_form
 
     def visit_callable_type(self, t: CallableType) -> bool:
         if self.ignore_in_type_obj and t.is_type_obj():