Make IntEnum members callable while keeping their comparison behavior

Question

Let's say the class Color is an enumeration:

from enum import IntEnum
class Color(IntEnum):
    red = 1
    blue = 2
    green = 3

Now, let's say I want to associate each of those member with a function AND keeping their value. i.e doing something like this

Colors.red()

Would call a function that I'd assign in someway. And keeping the IntEnum behavior, meaning that

Colors.red == 1

Would still return True

I guess I need to define a new superclass for my enumeration, but how?

The only thing near that that I have managed to do is:

class Colors(Enum):
    red = (1, f)
    blue = (2, g)
    green = (3, h)

Where f g and h are functions. However. With this method, the only way to access the value is

Colors.red.value[0]

And to call the function this is

Colors.red.value[1](args)

Wich is, in my opinion, pretty ugly, especially when names are a bit longs, when you have to call function a lot of times and with a lots of arguments

So is there a way to do what I want or do I need to stick to the uggly version?

Answer 1

Against my better judgement, I've decided to unravel the magic threads that hold Enums together. Here's what I came up with:

from enum import Enum

class FuncEnum(Enum):
    def __init__(self, val, func=None):
        self.val = val
        self.func = func or lambda: None
    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)
    def __eq__(self, other):
        return self.val == other
    def __ge__(self, other):
        return self.val >= other
    # and etc... not sure how @total_ordering will work here!

class Color(FuncEnum):
    red = (1, lambda x: x**2)
    blue = (2, lambda x: x**3)
    green = (3, lambda x: x**4)

Color.red(4)  # 16
Color.red == 1  # True
Color.blue(2)  # 8

This also supports reassignment, so afterwards:

Color.red.func = str.lower
Color.red("LOWERCASE")  # "lowercase"

Answer 2

Simple and plain, although a nightmare for readability and comprehensibility. Simply implements Python callable interface by providing a __call__ method. Bury it in some module, implement __all__ and never import anything besides Color class.

from enum import IntEnum
f1 = lambda x: x
f2 = lambda x: x*x
f3 = lambda x: x*x*x

class Color(IntEnum):
    red = 1
    blue = 2
    green = 3
    def __call__(self, *args, **kwargs):
        # alternatively some generic function of self.value
        return LUT[self](*args, **kwargs)


LUT = {Color.red: f1, Color.blue: f2, Color.green: f3}

assert Color.red(1) == 1
assert Color.blue(2) == 4
assert Color.green(3) == 27
assert Color.blue > Color.red

Edit: Much better (more readable, less hacky etc.) would be to define explicit callable interface on Enum.

class Color(IntEnum):
    red = 1
    blue = 2
    green = 3
    def get_something(self, x):
        # e.g.
        return x ** self.value

Any users of this Enum would do something like:

o = Obj()
o.color.get_something(123)

With with proper name of get_something would be readable, easy to understand for pylint and other static analysis tools and less error-prone.