How to define enum values that are functions?

Question

I have a situation where I need to enforce and give the user the option of one of a number of select functions, to be passed in as an argument to another function:

I really want to achieve something like the following:

from enum import Enum

#Trivial Function 1
def functionA():
    pass

#Trivial Function 2
def functionB():
    pass

#This is not allowed (as far as i can tell the values should be integers)
#But pseudocode for what I am after
class AvailableFunctions(Enum):
    OptionA = functionA
    OptionB = functionB

So the following can be executed:

def myUserFunction(theFunction = AvailableFunctions.OptionA):
   #Type Check
   assert isinstance(theFunction,AvailableFunctions) 

   #Execute the actual function held as value in the enum or equivalent
   return theFunction.value() 

Answer 1

It should be noted that as of Python 3.13.0 the partial version no longer works. This is due to a breaking change, as mentioned in this issue.

However, the wrapper class approach does work in 3.13.

The answer provided by STerliakov looks like it's the most natural one, but as mentioned it only works for Python >= 3.11.

Answer 2

I'm combining what I found with this and that answer and have come up with the following for Python 3.11+:

import enum
class A(enum.Enum):
    a = enum.member(func_a)
    b = enum.member(func_b)

Answer 3

Since Python 3.11 there is much more concise and understandable way. member and nonmember functions were added to enum among other improvements, so you can now do the following:

from enum import Enum, member, nonmember

def fn(x):
    print(x)

class MyEnum(Enum):
    x = nonmember(1)
    meth = fn
    mem = member(fn)
    @classmethod
    def this_is_a_method(cls):
        print('No, still not a member')
    def this_is_just_function():
        print('No, not a member')
    @member
    def this_is_a_member(x):
        print('Now a member!', x)

And now

>>> list(MyEnum)
[<MyEnum.mem: <function fn at ...>>, <MyEnum.this_is_a_member: <function MyEnum.this_is_a_member at ...>>]

>>> MyEnum.meth(1)
1

>>> MyEnum.mem(1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'MyEnum' object is not callable

>>> MyEnum.mem.value(1)
1

>>> MyEnum.this_is_a_method()
No, still not a member

>>> MyEnum.this_is_just_function()
No, not a member

>>> MyEnum.this_is_a_member()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'MyEnum' object is not callable

>>> MyEnum.this_is_a_member.value(1)
Now a member! 1

>>> MyEnum.x
1

Of course, you can add __call__ to avoid explicit .value calls:

class MyEnum(Enum):
    ...
    def __call__(self, *args, **kwargs):
        return self.value(*args, **kwargs)

and now

>>> MyEnum.mem(1)
1

Note that this cannot be typed correctly unless all your enum member functions share the same signature.

Answer 4

Just to add to this: there a way to make it work without partial or member in downstream code, but you have to go deep into metaclasses and the implementation of Enum. Building on the other answers, this works:

from enum import Enum, EnumType, _EnumDict, member
import inspect


class _ExtendedEnumType(EnumType):
    # Autowraps class-level functions/lambdas in enum with member, so they behave as one would expect
    # I.e. be a member with name and value instead of becoming a method
    # This is a hack, going deep into the internals of the enum class
    # and performing an open-heart surgery on it...
    def __new__(metacls, cls: str, bases, classdict: _EnumDict, *, boundary=None, _simple=False, **kwds):
        non_members = set(classdict).difference(classdict._member_names)
        for k in non_members:
            if not k.startswith("_"):
                if classdict[k].__class__ in [classmethod, staticmethod]:
                    continue
                # instance methods don't make sense for enums, and may break callable enums
                if "self" in inspect.signature(classdict[k]).parameters:
                    raise TypeError(
                        f"Instance methods are not allowed in enums but found method"
                        f" {classdict[k]} in {cls}"
                    )
                # After all the input validation, we can finally wrap the function
                # For python<3.11, one should use `functools.partial` instead of `member`
                callable_val = member(classdict[k])
                # We have to use del since _EnumDict doesn't allow overwriting
                del classdict[k]
                classdict[k] = callable_val
                classdict._member_names[k] = None
        return super().__new__(metacls, cls, bases, classdict, boundary=boundary, _simple=_simple, **kwds)

class ExtendedEnum(Enum, metaclass=_ExtendedEnumType):
    pass

Now you can do:

class A(ExtendedEnum):
    a = 3
    b = lambda: 4
    
    @classmethod
    def afunc(cls):
        return 5
    
    @staticmethod
    def bfunc():
        pass

Everything will work as expected.

PS: For some more Enum magic, I also like to add

    def __getitem__(cls, item):
        if hasattr(item, "name"):
            item = item.name
        # can't use [] because of particularities of super()
        return super().__getitem__(item)

to _ExtendedEnumType, so that A[A.a] works.

And I also recommend to make the callable enum, as proposed above.

`

Answer 5

Another less clunky solution is to put the functions in a tuple. As Bakuriu mentioned, you may want to make the enum callable.

from enum import Enum

def functionA():
    pass

def functionB():
    pass

class AvailableFunctions(Enum):
    OptionA = (functionA,)
    OptionB = (functionB,)

    def __call__(self, *args, **kwargs):
        self.value[0](*args, **kwargs)

Now you can use it like this:

AvailableFunctions.OptionA() # calls functionA

Answer 6

Building on top of @bakuriu's approach, I just want to highlight that we can also use dictionaries of multiple functions as values and have a broader polymorphism, similar to enums in Java. Here is a fictitious example to show what I mean:

from enum import Enum, unique

@unique
class MyEnum(Enum):
    test = {'execute': lambda o: o.test()}
    prod = {'execute': lambda o: o.prod()}

    def __getattr__(self, name):
        if name in self.__dict__:
            return self.__dict__[name]
        elif not name.startswith("_"):
            value = self.__dict__['_value_']
            return value[name]
        raise AttributeError(name)

class Executor:
    def __init__(self, mode: MyEnum):
        self.mode = mode

    def test(self):
        print('test run')

    def prod(self):
        print('prod run')

    def execute(self):
        self.mode.execute(self)

Executor(MyEnum.test).execute()
Executor(MyEnum.prod).execute()

Obviously, the dictionary approach provides no additional benefit when there is only a single function, so use this approach when there are multiple functions. Ensure that the keys are uniform across all values as otherwise, the usage won't be polymorphic.

The __getattr__ method is optional, it is only there for syntactic sugar (i.e., without it, mode.execute() would become mode.value['execute']().

Since dictionaries can't be made readonly, using namedtuple would be better and require only minor changes to the above.

from enum import Enum, unique
from collections import namedtuple

EnumType = namedtuple("EnumType", "execute")

@unique
class MyEnum(Enum):
    test = EnumType(lambda o: o.test())
    prod = EnumType(lambda o: o.prod())

    def __getattr__(self, name):
        if name in self.__dict__:
            return self.__dict__[name]
        elif not name.startswith("_"):
            value = self.__dict__['_value_']
            return getattr(value, name)
        raise AttributeError(name)

Answer 7

Your assumption is wrong. Values can be arbitrary, they are not limited to integers. From the documentation:

The examples above use integers for enumeration values. Using integers is short and handy (and provided by default by the Functional API), but not strictly enforced. In the vast majority of use-cases, one doesn’t care what the actual value of an enumeration is. But if the value is important, enumerations can have arbitrary values.

However the issue with functions is that they are considered to be method definitions instead of attributes!

In [1]: from enum import Enum

In [2]: def f(self, *args):
   ...:     pass
   ...: 

In [3]: class MyEnum(Enum):
   ...:     a = f
   ...:     def b(self, *args):
   ...:         print(self, args)
   ...:         

In [4]: list(MyEnum)  # it has no values
Out[4]: []

In [5]: MyEnum.a
Out[5]: <function __main__.f>

In [6]: MyEnum.b
Out[6]: <function __main__.MyEnum.b>

You can work around this by using a wrapper class or just functools.partial or (only in Python2) staticmethod:

from functools import partial

class MyEnum(Enum):
    OptionA = partial(functionA)
    OptionB = staticmethod(functionB)

Sample run:

In [7]: from functools import partial

In [8]: class MyEnum2(Enum):
   ...:     a = partial(f)
   ...:     def b(self, *args):
   ...:         print(self, args)
   ...:         

In [9]: list(MyEnum2)
Out[9]: [<MyEnum2.a: functools.partial(<function f at 0x7f4130f9aae8>)>]

In [10]: MyEnum2.a
Out[10]: <MyEnum2.a: functools.partial(<function f at 0x7f4130f9aae8>)>

Or using a wrapper class:

In [13]: class Wrapper:
    ...:     def __init__(self, f):
    ...:         self.f = f
    ...:     def __call__(self, *args, **kwargs):
    ...:         return self.f(*args, **kwargs)
    ...:     

In [14]: class MyEnum3(Enum):
    ...:     a = Wrapper(f)
    ...:     

In [15]: list(MyEnum3)
Out[15]: [<MyEnum3.a: <__main__.Wrapper object at 0x7f413075b358>>]

---

Also note that if you want you can define the __call__ method in your enumeration class to make the values callable:

In [1]: from enum import Enum

In [2]: def f(*args):
   ...:     print(args)
   ...:     

In [3]: class MyEnum(Enum):
   ...:     a = partial(f)
   ...:     def __call__(self, *args):
   ...:         self.value(*args)
   ...:         

In [5]: MyEnum.a(1,2,3)   # no need for MyEnum.a.value(1,2,3)
(1, 2, 3)

Answer 8

In addition to the answer of Bakuriu... If you use the wrapper approach like above you loose information about the original function like __name__, __repr__ and so on after wrapping it. This will cause problems for example if you want to use sphinx for generation of source code documentation. Therefore add the following to your wrapper class.

class wrapper:
    def __init__(self, function):
        self.function = function
        functools.update_wrapper(self, function)
    def __call__(self,*args, **kwargs):
        return self.function(*args, **kwargs)
    def __repr__(self):
        return self.function.__repr__()