Python Patterns: Enum
Things often come in sets, for example, States, Pokémon, Playing cards, etc.
Each set has items that belong to them (like California, Charizard, Jack of Clubs) and checking if an item is a valid member is a common task. Some collections (playing cards, for example) are also orderable; twos come before fives which come before Kings.
There are many ways to represent members from these sets in Python:
- Unique string:
class Pokemon: ...
- Tuples (or namedtuples):
"PR" a valid state? Is
Pokemon("Digimon") a member of Pokemon? Is
("Lotus", "Black") a playing card? We could keep a separate
set() of all valid members to check, but then we have to maintain it.
Thankfully, Python provides a way to create these sets and their members at the same time: enumerations, or enums.
Without using enums we might implement a standard playing card as follows:
This class works with the standard comparison operators, but to do so we had to write a bit of an annoying
__rank_to_value() function; otherwise Aces and Kings would be tough to compare to 2s and 3s!
With that done, we can now declare cards easily enough:
Did you catch all the errors? We could write some error checking in the class, but it would again be a bit tedious. Instead, let’s implement this using enums. Enums will let us represent the suits and ranks, check that they are valid, and order based on value, without writing a lot of extra code.
Playing Cards with Enums
An enum has exactly the properties we want:
- We can test membership, so only real suits and ranks are allowed.
- We can order them, so we know that King > Jack > Ten.
First, we define the suits:
auto() sets the values and insures that they are unique. The members are not orderable (so
CardSuit.CLUBS > CardSuit.DIAMONDS will raise an error), but do have equality (so
CardSuit.CLUBS != CardSuit.DIAMONDS works). We can also test membership easily, allowing us to ensure only valid suits are accepted.
An example of some of the properties:
Second, we define a
CardValue, this time using
IntEnum because we want the values to be comparable.
IntEnums are orderable so
CardRank.TEN < CardRank.KING. The decorator
@unique adds a check that makes sure we haven’t double assigned any values.
Now the card class is easy to implement:
It is now much easier to catch errors in our card definitions:
Not only are they obvious by eye (
"Stars" is clearly not a
CardSuit), but the runtime will even raise an error!