Subclass Enum to add validation

Question

The Python docs for the enum module contains the following example of subclassing Enum. The resulting class can be used to create enums that also validate that they have no two members with the same value.

>>> class DuplicateFreeEnum(Enum):
...     def __init__(self, *args):
...         cls = self.__class__
...         if any(self.value == e.value for e in cls):
...             a = self.name
...             e = cls(self.value).name
...             raise ValueError(
...                 "aliases not allowed in DuplicateFreeEnum:  %r --> %r"
...                 % (a, e))

However, as a method of adding validation, this method is inelegant and restrictive. __init__ is called once for each member, whereas in order to validate an enum as a whole, it makes more sense to look at every member of the enum together.

For instance, how would I validate that an enum has precisely two members, as below?

class PreciselyTwoEnum(Enum):
    ...  # ???

class Allowed(PreciselyTwoEnum):
    FOO = 1
    BAR = 2

class Disallowed(PreciselyTwoEnum):  # Should raise an error
    BAZ = 3

Can this be accomplished with a clever implementation of __init__? Is there another method that could be used — perhaps one that is called on the enum after it has been fully created?

score 3 · Answer 1 · answered Aug 25 '22 at 20:07

3

I don't know how to subclass the Enum to achieve the required functionality, but it can be easily done with a class decorator:

from enum import Enum

def PreciselyTwoEnum(cls):
    members = len(cls.__members__)
    if members != 2:
        raise ValueError(f"two members expected in {cls.__name__!r}, but got {members}")
    return cls 

@PreciselyTwoEnum
class Allowed(Enum):
    FOO = 1 
    BAR = 2 

@PreciselyTwoEnum
class Disallowed(Enum):  # Should raise an error
    BAZ = 3

answered Aug 25 '22 at 20:07

VPfB

14,927
6
41
75

1

+1 upcoming in 3.11 are features which also do similar tasks using class decorator - see [`@enum.verify`](https://docs.python.org/3.11/library/enum.html#enum.verify) – wim Aug 25 '22 at 20:23

score 2 · Accepted Answer · answered Aug 25 '22 at 20:35

__init_subclass__ is what you are looking for¹:

class PreciselyTwoEnum(Enum):
    def __init_subclass__(cls):
        if len(cls.__members__) != 2:
            raise TypeError("only two members allowed")

and in use:

>>> class Allowed(PreciselyTwoEnum):
...     FOO = 1 
...     BAR = 2 
... 

>>> class Disallowed(PreciselyTwoEnum):  # Should raise an error
...     BAZ = 3 
... 

Traceback (most recent call last):
  ...
TypeError: only two members allowed

[1] __init_subclass__ for Enum only works correctly in Python 3.11 or later, or by using the external aenum library².

[2] Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Subclass Enum to add validation

2 Answers2