CODEWITHSUNDEEP
pythonenums
Martin Akbaev
Martin Akbaev

Reputation: 37

How to get new Enum with members as enum using EnumMeta Python 3.6

I'm trying translate some code from python 2 to python 3 and had trouble with this code:

from enum import Enum, EnumMeta


class KeyCode(Enum):
    A = 1
    B = 2
    C = 3


class KeyMeta(EnumMeta):
    def __new__(mcs, name, bases, dct):
        dct.update({e.name: e.value for e in KeyCode})
        return super(KeyMeta, mcs).__new__(mcs, name, bases, dct)


class Key(Enum, metaclass=KeyMeta):
    pass


print(repr(Key.A))

The problem is that for python2 Key.A is enum < Key.A: 1>, whereas for python3 it is < class 'int'>.

Upvotes: 1

Views: 738

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121904

In Python 3, the dct object is not a regular dictionary, but a subclass dedicated to helping create enums, set via the __prepare__ attribute of the metaclass. It's update() method, however, is not altered from the base dictionary.

It normally expects members to be set via it's __setitem__ method (e.g. dct[name] = value); do so in your __new__ method too:

class KeyMeta(EnumMeta):
    def __new__(mcs, name, bases, dct):
        for e in KeyCode:
            dct[e.name] = e.value
        return super(KeyMeta, mcs).__new__(mcs, name, bases, dct)

Now each of the names A, B and C are recorded as being enum member names and the EnumMeta class takes it from there. Without using the specialised __setitem__ implementation, the names are seen as some other type of attribute instead.

With the above change you get the expected output:

>>> class Key(Enum, metaclass=KeyMeta):
...     pass
...
>>> Key.A
<Key.A: 1>

The same for e in KeyCode: loop will continue to work in Python 2 as well.

Upvotes: 1

Related Questions