Search code examples
pythonenumsintegeroverriding

Customize "return by value" in IntEnum


I have a class that extends IntEnum class that defines positions in a bit-encoded variable:

from enum import IntEnum

class Bits(IntEnum):

    @classmethod
    def data(cls, value: int):
        return [ e for e in cls if (1 << e) & value ]

class Status(Bits):
    READY = 0
    ERROR = 1
    WARNING = 2

x = Status.data(3) # x <- [<Status.READY: 0>, <Status.ERROR: 1>]
y = Status.data(4) # y <- [<Status.WARNING: 2>]
z = Status.data(8) # z <- []

Would it be possible to customize "return by value" of IntEnum without breaking anything? Since there will be multiple classes that will extend the Bits class, ideally this functionality should be implemented in the Bits class.

What I want to achieve is this:

x = Status(3) # x <- [<Status.READY: 0>, <Status.ERROR: 1>]
y = Status(4) # y <- [<Status.WARNING: 2>]
z = Status(8) # z <- []

I tried overriding the __call__ method (see enum — Support for enumerations), but it does not seem to be called in this case.


Solution

  • Yes you could customize the return value of an IntEnum class by defining a custom __new__ method for your Status class, this way it will change the behavior of how new instances are created without breaking the functionality of the IntEnum.

    from enum import IntEnum, EnumMeta
    
    class BitsMeta(EnumMeta):
        def __call__(cls, value):
            if isinstance(value, int):
                return [e for e in cls if (1 << e.value) & value]
            return super().__call__(value)
    
    class Bits(IntEnum, metaclass=BitsMeta):
    
        @classmethod
        def data(cls, value: int):
            return [e for e in cls if (1 << e.value) & value]
    
    class Status(Bits):
        READY = 0
        ERROR = 1
        WARNING = 2
    
    x = Status(3)  # x <- [<Status.READY: 0>, <Status.ERROR: 1>]
    y = Status(4)  # y <- [<Status.WARNING: 2>]
    z = Status(8)  # z <- []
    
    print(x)
    print(y)
    print(z)