Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
20 changes: 19 additions & 1 deletion mypy/typeops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
57 changes: 57 additions & 0 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading