Search code examples
pythonenumsflags

Flag usage for multibranch logic


I need to do something like this:

from enum import Flag, auto

class WISENESS(Flag):
    Y = auto()
    M = auto()
    D = auto()
    YM = Y | M
    YD = Y | D
    MD = M | D
    YMD = Y | M | D

first_case = WISENESS.Y

first_case == WISENESS.Y # True
first_case == WISENESS.M # False
first_case == WISENESS.D # False

###

second_case = WISENESS.YD

second_case == WISENESS.Y # True
second_case == WISENESS.M # False
second_case == WISENESS.D # True

####

third_case = WISENESS.YMD

third_case == WISENESS.Y # True
third_case == WISENESS.M # True
third_case == WISENESS.D # True

I.e. depending on the flag value it will be true in some cases. For example, I may need to perform an operation for all possible cases, or only for two of them. Like this example here:

if WISENESS.Y:
    do_something_in_case_of_Y_or_MY_or_YD_or_YMD()
if WISENESS.M:
    do_something_in_case_of_M_or_MD_or_YM_or_YMD()
if WISENESS.D:
    do_something_in_case_of_D_or_MD_or_YD_or_YMD()

I tried to use Flag from the enum module in the standard library, guessing it could help me in this case, but either I don't understand how it works, or I must achieve my goal in a different way.


Solution

  • The built-in way to check for Flag membership is with the standard Python in operator:

    >>> second_case in WISENESS.Y
    True
    

    and your final example would be:

    some_flag = ...
    
    if WISENESS.Y in some_flag:
        do_something_in_case_of_Y_or_MY_or_YD_or_YMD()
    if WISENESS.M in some flag:
        do_something_in_case_of_M_or_MD_or_YM_or_YMD()
    if WISENESS.D in some flag:
        do_something_in_case_of_D_or_MD_or_YD_or_YMD()