Do not narrow Flag enums by member equality

narrow_type_by_identity_equality expands enum instances into a union of
member literals and rules out the compared-away member in the negative
branch. For enum.Flag subclasses this is unsound: a Flag value can hold
any combination of members, so the declared members are not an exhaustive
value set. After `programs != Programs.P2`, programs can still be
`Programs.P1 | Programs.P2`, so narrowing away the P2 member rejects
`Programs.P2 in programs` -- valid at runtime -- against typeshed's
`Flag.__contains__(self, other: Self)` (gh-21937).

Skip equality/identity narrowing when either operand is a Flag value;
the sound type for such a value is the enum class itself. Regular enum
narrowing is unchanged, and existing Flag bitwise-operator behavior is
untouched.

Author: istoolsfox

mypy/checker.py

@@ -246,6 +246,7 @@ def __init__(self) -> None:
     false_only,
     fixup_partial_type,
     function_type,
+    is_flag_enum_type,
     is_literal_type_like,
     is_singleton_equality_type,
     is_singleton_identity_type,
@@ -7075,6 +7076,15 @@ def narrow_type_by_identity_equality(
                 if should_coerce_literals:
                     target_type = coerce_to_literal(target_type)
 
+                # A Flag value can hold any combination of members, so the declared
+                # members are not an exhaustive set of values: ruling out e.g.
+                # `Programs.P2` does not rule out `Programs.P1 | Programs.P2`. Narrowing
+                # by member equality is therefore unsound for Flag enums.
+                if is_flag_enum_type(get_proper_type(expr_type)) or is_flag_enum_type(
+                    get_proper_type(target_type)
+                ):
+                    continue
+
                 # Morally what we want to do is narrow for each branch based on:
                 # `if_type, else_type = conditional_types(expr_type, target)`
                 # What we actually do is first munge expr_type based on target_type to handle some

mypy/typeops.py

@@ -9,7 +9,7 @@
 
 import itertools
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, TypeVar, cast
+from typing import Any, Final, TypeVar, cast
 
 from mypy.checker_state import checker_state
 from mypy.copytype import copy_type
@@ -1076,6 +1076,24 @@ def is_singleton_equality_type(typ: ProperType) -> bool:
     return isinstance(typ, LiteralType) or is_singleton_identity_type(typ)
 
 
+FLAG_ENUM_BASES: Final = ("enum.Flag", "enum.IntFlag")
+
+
+def is_flag_enum_type(typ: ProperType) -> bool:
+    """Is this type an instance of, or a member of, an enum.Flag subclass?
+
+    A Flag value can hold any combination of members, so the declared members
+    are not an exhaustive enumeration of the possible values.
+    """
+    if isinstance(typ, LiteralType) and typ.is_enum_literal():
+        typ = typ.fallback
+    return isinstance(typ, Instance) and typ.type.is_enum and is_flag_enum_class(typ.type)
+
+
+def is_flag_enum_class(info: TypeInfo) -> bool:
+    return any(base.fullname in FLAG_ENUM_BASES for base in info.mro)
+
+
 def try_expanding_sum_type_to_union(typ: Type, target_fullname: str | None) -> Type:
     """Attempts to recursively expand any enum Instances with the given target_fullname
     into a Union of all of its component LiteralTypes.

test-data/unit/check-enum.test

@@ -485,6 +485,63 @@ if int():
     x = x | C.b
 [builtins fixtures/enum.pyi]
 
+[case testFlagEnumEqualityNarrowing]
+from enum import Flag, auto
+
+class Programs(Flag):
+    NONE = 0
+    P1 = auto()
+    P2 = auto()
+    P3 = auto()
+    ALL = P1 | P2 | P3
+
+def f1(programs: Programs) -> None:
+    if programs == Programs.NONE:
+        return
+    if programs == Programs.P2:
+        return
+    # Flag values can hold any combination of members, so narrowing by member
+    # equality would be unsound here (e.g. `Programs.P1 | Programs.P2`).
+    reveal_type(programs)  # N: Revealed type is "__main__.Programs"
+
+def f2(programs: Programs) -> None:
+    if programs != Programs.P2:
+        reveal_type(programs)  # N: Revealed type is "__main__.Programs"
+[builtins fixtures/primitives.pyi]
+
+[case testFlagEnumIdentityNarrowing]
+from enum import Flag, auto
+
+class Programs(Flag):
+    NONE = 0
+    P1 = auto()
+    P2 = auto()
+    P3 = auto()
+    ALL = P1 | P2 | P3
+
+def f(programs: Programs) -> None:
+    if programs is not Programs.P2:
+        reveal_type(programs)  # N: Revealed type is "__main__.Programs"
+[builtins fixtures/primitives.pyi]
+
+[case testIntFlagEnumEqualityNarrowing]
+from enum import IntFlag, auto
+
+class Programs(IntFlag):
+    NONE = 0
+    P1 = auto()
+    P2 = auto()
+    P3 = auto()
+    ALL = P1 | P2 | P3
+
+def f(programs: Programs) -> None:
+    if programs == Programs.NONE:
+        return
+    if programs == Programs.P2:
+        return
+    reveal_type(programs)  # N: Revealed type is "__main__.Programs"
+[builtins fixtures/primitives.pyi]
+
 [case testAnonymousEnum]
 from enum import Enum
 class A: