C. Loew
C. Loew

Reputation: 303

Python 3 Enums with Function Values

I noticed an oddity in the Python 3 Enums (link).
If you set the value of an Enum to a function, it prevents the attribute from being wrapped as an Enum object, which prevents you from being able to use the cool features like EnumCls['AttrName'] to dynamically load the attribute.

Is this a bug? Done on purpose?
I searched for a while but found no mention of restricted values that you can use in an Enum.

Here is sample code that displays the issue:

class Color(Enum):
    Red = lambda: print('In Red')
    Blue = lambda: print('In Blue')

print(Color.Red)    # <function> - should be Color.Red via Docs
print(Color.Blue)   # <function> - should be Color.Bluevia Docs
print(Color['Red']) # throws KeyError - should be Color.Red via Docs

Also, this is my first time asking, so let me know if there's anything I should be doing differently! And thanks for the help!

Upvotes: 25

Views: 14412

Answers (7)

Michael Panchenko
Michael Panchenko

Reputation: 461

There is a way to make it work without partial, FunctionProxy or anything like that, 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.

See also this thread.

`

Upvotes: 0

Darkdragon84
Darkdragon84

Reputation: 911

You can also use functools.partial to trick the enum into not considering your function a method of Color:

from functools import partial
from enum import Enum

class Color(Enum):
    Red = partial(lambda: print('In Red'))
    Blue = partial(lambda: print('In Blue'))

With this you can access name and value as expected.

Color.Red
Out[17]: <Color.Red: functools.partial(<function Color.<lambda> at 0x7f84ad6303a0>)>
Color.Red.name
Out[18]: 'Red'
Color.Red.value()
In Red

Upvotes: 6

K3---rnc
K3---rnc

Reputation: 7059

You can override the __call__ method:

from enum import Enum, auto

class Color(Enum):
    red = auto()
    blue = auto()

    def __call__(self, *args, **kwargs):
        return f'<font color={self.name}>{args[0]}</font>'

Can then be used:

>>> Color.red('flowers')
<font color=red>flowers</font>

Upvotes: 13

Razzle Shazl
Razzle Shazl

Reputation: 1330

Initially, I thought your issue was just missing commas because I got the output you were expecting.:

from enum import Enum

class Color(Enum):
    Red = lambda: print('In Red'),
    Blue = lambda: print('In Blue'),

print(Color.Red)
print(Color.Blue)
print(Color['Red'])

output (python3.7)

$ /usr/local/opt/python/bin/python3.7 ~/test_enum.py
Color.Red
Color.Blue
Color.Red

@BernBarn was kind enough to explain that in my solution that a tuple is being created, and to invoke the function would require dereferencing value[0]. There is already another answer using value[0] in this way. I miss rb for this.

Upvotes: -1

smarie
smarie

Reputation: 5256

I ran into this issue recently, found this post, and first was tempted to use the wrapper pattern suggested in the other related post. However eventually I found out that this was a bit overkill for what I had to do. In the past years this happened to me several times with Enum, so I would like to share this simple experience feedback:

if you need an enumeration, ask yourself whether you actually need an enum or just a namespace.

The difference is simple: Enum members are instances of their host enum class, while namespace members are completely independent from the class, they are just located inside.

Here is an example of namespace containing callables, with a get method to return any of them by name.

class Foo(object):
    """ A simple namespace class with a `get` method to access members """
    @classmethod
    def get(cls, member_name: str):
        """Get a member by name"""
        if not member_name.startswith('__') and member_name != 'get':
            try:
                return getattr(cls, member_name)
            except AttributeError:
                pass
        raise ValueError("Unknown %r member: %r" % (cls.__name__, member_name))

    # -- the "members" --

    a = 1

    @staticmethod
    def welcome(name):
        return "greetings, %s!" % name

    @staticmethod
    def wave(name):
        return "(silently waving, %s)" % name


w = Foo.get('welcome')
a = Foo.get('a')
Foo.get('unknown')  # ValueError: Unknown 'Foo' member: 'unknown'

See also this post on namespaces.

Upvotes: 0

Ceppo93
Ceppo93

Reputation: 1046

If someone need/want to use Enum with functions as values, its possible to do so by using a callable object as a proxy, something like this:

class FunctionProxy:
    """Allow to mask a function as an Object."""
    def __init__(self, function):
        self.function = function

    def __call__(self, *args, **kwargs):
        return self.function(*args, **kwargs)

A simple test:

from enum import Enum
class Functions(Enum):
    Print_Function = FunctionProxy(lambda *a: print(*a))
    Split_Function = FunctionProxy(lambda s, d='.': s.split(d))

Functions.Print_Function.value('Hello World!')
# Hello World!
Functions.Split_Function.value('Hello.World.!')
# ['Hello', 'World', '!']

Upvotes: 6

BrenBarn
BrenBarn

Reputation: 251598

The documentation says:

The rules for what is allowed are as follows: _sunder_ names (starting and ending with a single underscore) are reserved by enum and cannot be used; all other attributes defined within an enumeration will become members of this enumeration, with the exception of __dunder__ names and descriptors (methods are also descriptors).

A "method" is just a function defined inside a class body. It doesn't matter whether you define it with lambda or def. So your example is the same as:

class Color(Enum):
    def Red():
        print('In Red')
    def Blue():
        print('In Blue')

In other words, your purported enum values are actually methods, and so won't become members of the Enum.

Upvotes: 10

Related Questions