en fr

Use metaclasses to create simple enums in Python 3

Posted on 2016-08-01 in Programmation Last modified on: 2016-08-15

Since version 3.4, Python has an Enum class. It allows you to create enums with some neat properties:

  • Iteration.
  • Length.
  • Access to the members with the "attribute" or "dict" notation.

I leave you to the documentation for the details.

In my case, I am using them to track some values of requests send to my application (eg: the name of the request or the state of a slot – opened, closed, taken). In this use case, I have some problem with the "standard" enums as provided by the enum package:

  • I have to enter the name and value for each member which is kind of sad since they are the same:

    class SlotState(Enum):
    OPEN = 'OPEN'
    CLOSED = 'CLOSED'
    RESERVED = 'RESERVED'
    TAKEN = 'TAKEN'

  • I always forget the .value when I am testing: I do SlotState.OPEN == request['slot']['state'] instead of SlotState.OPEN.value == request['slot']['state'].

So I asked myself: can I create my own enums that would solve these problems? I want them to:

  • Use the name of the attribute to fill its value.
  • Test directly with the attribute: MonEnum.membre == 'VALEUR'.
  • Not to require to use caps in all access, ie access with either MonEnum['Valeur'] or MonEnum['VALEUR']
  • Keep the neat properties of "true" Python enums:
    • Length
    • Loop over the elements
    • Access to the elements with either the "attribute" or the dict notation.
    • Test if an element is part of an enum with the in keyword.

The answer is, it is possible with metaclasses and by getting some inspiration from EnumMeta the metaclass of the Enum class.

What is a metaclass? You probably know that in Python everything is an object. Well, classes are no exceptions: they are instances of a metaclass. This allows you to make operations on a class instead of an object. For instance, if you declare a __getitem__ method on a class, you must use it on an instance of the class like this instance['toto'] but you cannot use it directly on the class like this: MyClass['toto'].

However, you can do this with a metaclass: if you declare __getitem__ in the metaclass of MyClass you can use MyClass['toto'].

Here's my SimpleEnumMeta metaclass (explanations below):

 1 class SimpleEnumMeta(type):
 2     def __new__(metacls, cls, bases, classdict):
 3          object_attrs = set(dir(type(cls, (object,), {})))
 4          simple_enum_cls = super().__new__(metacls, cls, bases, classdict)
 5          simple_enum_cls._member_names_ = set(classdict.keys()) - object_attrs
 6          non_members = set()
 7          for attr in simple_enum_cls._member_names_:
 8              if attr.startswith('_') and attr.endswith('_'):
 9                  non_members.add(attr)
10              else:
11                  setattr(simple_enum_cls, attr, attr)
12 
13          simple_enum_cls._member_names_.difference_update(non_members)
14 
15          return simple_enum_cls
16 
17     def __getitem__(cls, key):
18        return getattr(cls, key.upper())
19 
20     def __iter__(cls):
21        return (name for name in cls._member_names_)
22 
23     def __len__(cls):
24        return len(cls._member_names_)

How does it work?

  • The constructor __new__ (not to be mistaken with the initialization function __init__) allows us to interact with the class before its "creation". That allows us to affect values to the attribute of our class. To do that:
    1. We get the set of attributes common to all objects in order to exclude them from the set of attributes whose value we want to set automatically (line 3).
    2. We create the class thanks to the constructor of type the super class of all metaclasses (line 4).
    3. We store in a private property the set of attributes specific to our enum (line 5). This allows us to give the correct number of elements of you enum and to create an iterator to loop over them (methods __len__ and __iter__).
    4. For each attribute in the enum, we affect its value as the name of the attribute (line 7 to 11).
    5. We then remove the private attributes from the list of the list of attributes.
    6. In the end, we return the newly created class (ligne 15).
  • I don't think I need to explain the other methods.

To use it, you just need to create a class with the metaclass=SimpleEnumMeta parameter:

class SlotState(metaclass=SimpleEnumMeta):
    OPEN = ()
    CLOSED = ()
    RESERVED = ()
    TAKEN = ()
    AI = ()

You must assign a value to the attributes since:

class MyClass:
   test

is not valid Python code and results in a NameError exception. Using an empty tuple looks like a good way to write valid code while signaling you don't need to set the properties.

You can then check that your class behave as intended:

  • SlotState.OPEN == 'OPEN'
  • SlotState['open'] == 'OPEN'
  • len(SlotState) == 5
  • 'OPEN' in SlotState
  • for state in SlotState: print(state)

Et voilà ! If you have a question, please leave a comment.

Python