diff --git a/mypy/checker.py b/mypy/checker.py index 33ed5387554d..42268e68317c 100644 --- a/mypy/checker.py +++ b/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 diff --git a/mypy/typeops.py b/mypy/typeops.py index 8453da4dd31c..a8202d23fb7a 100644 --- a/mypy/typeops.py +++ b/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. diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 7c24f995f442..8c7346eabdfc 100644 --- a/test-data/unit/check-enum.test +++ b/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: