CODEWITHSUNDEEP
pythonenums
Snuh
Snuh

Reputation: 360

Returning default members when Enum member does not exist

I have an Enum where I would like a default member to be returned when a member does not exist inside of it. For example:

class MyEnum(enum.Enum):
    A = 12
    B = 24
    CUSTOM = 1


print(MyEnum.UNKNOWN) # Should print MyEnum.CUSTOM

I know I can use a metaclass like so:

class MyMeta(enum.EnumMeta):
    def __getitem__(cls, name):
        try:
            return super().__getitem__(name)
        except KeyError as error:
            return cls.CUSTOM

class MyEnum(enum.Enum,metaclass=MyMeta):
    ...

But that appears to only work if I access the Enum using MyEnum['UNKNOWN']. Is there a way that covers both methods of accessing members of an enum when the member doesn't exist?

Upvotes: 6

Views: 921

Answers (2)

Sergey Nudnov
Sergey Nudnov

Reputation: 1443

Compiling the answer of @BrokenBenchmark, and this answer, and a bit of knowledge, let's make default enum member a parameter:

class EnumWithDefaultMeta(enum.EnumMeta):

    @classmethod
    def __prepare__(metacls, name, bases, **kwargs):
        return super().__prepare__(name, bases)

    def __new__(metacls, name, bases, namespace, **kwargs):
        newclass = super().__new__(metacls, name, bases, namespace)
        newclass._missing_ = classmethod(metacls._missing_)
        return newclass

    def __init__(cls, name, bases, namespace, default='UNKNOWN'):
        cls.default = default
        super().__init__(name, bases, namespace)

    def __getitem__(cls, name):
        try:
            return super().__getitem__(name)
        except KeyError:
            return super().__getitem__(cls.default)

    def __getattr__(cls, name):
        try:
            # noinspection PyUnresolvedReferences
            return super().__getattr__(name)
        except AttributeError:
            # noinspection PyUnresolvedReferences
            return super().__getattr__(cls.default)

    def _missing_(cls, _value):
        return super().__getitem__(cls.default)

class MyEnum(enum.Enum, metaclass=EnumWithDefaultMeta, default='CUSTOM'):
    A = 12
    B = 24
    CUSTOM = 1

Output:

MyEnum(123)
>>> <MyEnum.CUSTOM: 1>
MyEnum.Test
>>> <MyEnum.CUSTOM: 1>
MyEnum['Test']
>>> <MyEnum.CUSTOM: 1>

Upvotes: 0

BrokenBenchmark
BrokenBenchmark

Reputation: 19252

Add a definition for __getattr__ to the metaclass:

class MyMeta(enum.EnumMeta):
    def __getitem__(cls, name):
        try:
            return super().__getitem__(name)
        except KeyError as error:
            return cls.CUSTOM
    
    def __getattr__(cls, name):
        try:
            return super().__getattr__(name)
        except AttributeError as error:
            return cls.CUSTOM

Then, your code will output:

MyEnum.CUSTOM

Upvotes: 4

Related Questions