fixed Incorrect type of enum in 'if' clause #20234

Author: sherlock2215

mypy/checker.py

@@ -241,6 +241,7 @@
 from mypy.typevars import fill_typevars, fill_typevars_with_any, has_no_typevars
 from mypy.util import is_dunder, is_sunder
 from mypy.visitor import NodeVisitor
+from mypy.types import LiteralType
 
 T = TypeVar("T")
 
@@ -6517,7 +6518,6 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
         # Step 1: Obtain the types of each operand and whether or not we can
         # narrow their types. (For example, we shouldn't try narrowing the
         # types of literal string or enum expressions).
-
         operands = [collapse_walrus(x) for x in node.operands]
         operand_types = []
         narrowable_operand_index_to_hash = {}
@@ -6581,18 +6581,39 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
 
                 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 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_map[operands[left_index]] = remove_optional(item_type)
+                literal_types = []
+                if isinstance(get_proper_type(iterable_type), TupleType):
+                    # Check if this is an enum instance that can be narrowed
+                    tuple_type = get_proper_type(iterable_type)
+                    for i, item_type in enumerate(tuple_type.items):
+                        if isinstance(item_type, Instance):
+
+                            if item_type.type.is_enum:
+                                # Enum values in tuples are represented as Instance types, not LiteralType
+                                if hasattr(item_type, 'last_known_value') and item_type.last_known_value:
+                                    # Use the existing literal representation
+                                    literal_types.append(item_type.last_known_value)
+                                else:
+                                    # using the instance directly
+                                    literal_types.append(item_type)
+                    # If we found enum literals in the tuple, narrow the left operand
+                    if literal_types:
+                        union_type = make_simplified_union(literal_types)
+                        # Applying type narrowing for the true branch of the 'in' check
+                        if_map[operands[left_index]] = union_type
 
                 if right_index in narrowable_operand_index_to_hash:
                     if_type, else_type = self.conditional_types_for_iterable(