fr en

Utiliser des métaclasses pour créer simplement des enums en Python 3

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

Depuis la version 3.4, Python dispose d'une classe Enum qui permet de créer des enums avec quelques propriétés intéressantes (itération, nombre d'éléments, accès aux éléments de l'enum comme ceux d'un objet ou d'un dictionnaire). Je vous laisse lire la documentation pour les détails.

Cependant, dans mon cas, je les utilise pour traquer certaines valeurs d'une requête faites sur une application (par exemple : le nom de la requête envoyée au serveur ou l'état d'un créneau – ouvert, fermé, pris). Les enum de base me posent plusieurs problèmes :

  • Il faut entrer le nom et la valeur pour chaque membre, comme ceci :

    class SlotState(Enum):
        OPEN = 'OPEN'
        CLOSED = 'CLOSED'
        RESERVED = 'RESERVED'
        TAKEN = 'TAKEN'
    
  • J'oublie le .value à chaque fois lorsque je fais un test entre la requête et l'enum: je fais SlotState.OPEN == request['slot']['state'] au lieu de SlotState.OPEN.value == request['slot']['state'].

Du coup, je me suis demandé si on ne pouvais pas faire plus simple pour mon utilisation. Ce que j'aimerais c'est :

  • Utiliser le nom de l'attribut pour remplir automatiquement sa valeur.
  • Tester directement avec l'attribut MonEnum.membre == 'VALEUR'.
  • Ne pas avoir à mettre manuellement en majuscule les propriétés, ie pouvoir accéder aux éléments avec MonEnum['Valeur'] ou avec MonEnum['VALEUR'].
  • Garder les propriétés sympathiques des « vrais » enums Python :
    • Longueur.
    • Boucler sur tous les éléments de l'enum.
    • Accès aux membres avec la notation crochet et comme aux attributs d'une classe.
    • Tester si un élément existe dans l'enum avec le mot clé in.

C'est possible en utilisant les métaclasses et en s'inspirant un peu de EnumMeta la métaclasse de la classe Enum.

Qu'est-ce qu'une métaclasse ? Pour faire simple, vous savez certainement qu'en Python tout est objet et les classes ne font pas exception : ce sont des instances d'une métaclasse. Cela nous permet donc de faire des opérations sur la classe et non sur l'objet. Par exemple, si vous déclarer une méthode __getitem__ dans une classe, vous pouvez utiliser sur l'instance comme ceci instance['toto'] mais pas directement sur la classe comme cela MaClasse['toto'].

C'est ce que permet de faire une métaclasse : vous déclarer la méthode __getitem__ dans la métaclasse puis vous pouvez utiliser MaClasse['toto'].

Voici ma métaclasse SimpleEnumMeta (voir les explications ci-dessous) :

 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_)

Concernant le fonctionnement de cette métaclasse :

  • Le constructeur __new__ (à ne pas confondre avec la méthode d'initialisation __init__) nous permet d'interagir avec la classe avant sa "création" ce qui nous permet d'affecter des valeurs à nos attributs de classe. Pour cela :
    1. On récupère les attributs communs à tous les objets pour pouvoir les exclure de la liste des attributs dont on va affecter la valeur automatiquement (ligne 3).
    2. On crée la classe grâce au constructeur de type, la classe mère de toutes les métaclasses (ligne 4).
    3. On stocke dans une propriété privée l'ensemble des attributs spécifiques à notre enum (ligne 5). Cela nous permet de donner le nombre d'élément contenu dans l'enum ainsi que de fournir un itérateur sur ces éléments (méthodes __len__ et __iter__).
    4. Pour chaque attribut de l'enum, on affecte sa valeur comme le nom de l'attribut (lignes 7 à 11).
    5. On supprime de la liste les attributs les éventuels attributs privés.
    6. Enfin, on renvoie la classe nouvellement créée (ligne 15).
  • Je pense que les autres méthodes parles d'elles mêmes

Pour l'utiliser, il suffit de créer une classe avec un argument metaclass=SimpleEnumMeta :

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

Vous devez assigner une valeur aux attributs puisque :

class MyClass:
   test

n'est pas du Python valide et provoque une exception de type NameError. L'utilisation d'un tuple vide me semble un bon compromis pour avoir du code valide et signaler que la valeur de l'attribut n'a pas d'importance.

On peut ensuite vérifier que l'on a bien les propriétés attendues :

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

Et voilà ! Si vous avez une question ou une remarque, n'hésiter pas à utiliser les commentaires.