CODEWITHSUNDEEP
pythonenums
Anonymous12358
Anonymous12358

Reputation: 489

Subclass Enum to add validation

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?

Upvotes: 4

Views: 567

Answers (2)

Ethan Furman
Ethan Furman

Reputation: 69248

__init_subclass__ is what you are looking for1:

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 v3.0 or later.2.

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

Upvotes: 2

VPfB
VPfB

Reputation: 17362

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

Upvotes: 3

Related Questions