Enums with fields in Python

Why this matters

Many programs have a small set of choices that never changes. A coffee comes in three sizes. A shop accepts four payment types. A card is a heart, a diamond, a club or a spade. Programmers call such a set of named choices an enum, short for enumeration, which means a list of named values.

A plain list of names is only half the story. A coffee size is not just the word "MEDIUM". It has a volume, a price and a short code like "M" that goes on a receipt. If you store those facts in separate variables, or in a pile of if statements, the code grows messy and easy to get wrong. An enum that carries its own data keeps the name and the facts together.

This lesson shows how to do that with enum.Enum in Python. An enum member can have attributes and methods, like any object. The lesson also covers lookups from a text code to a member, the checks Python makes, and the mistakes that come up most often.

You need to know what a class and an __init__ method are. The earlier lessons in this module cover them.

A plain enum first

The simplest enum is only a list of names. Import Enum from the enum module, and inherit from it.

The output is:

Day.MONDAY
MONDAY 1
True

Each name in the class body, such as MONDAY, is a member. A member is one fixed value of the enum, and Python creates it once, when the class is defined. You refer to it with the class name and a dot, such as Day.MONDAY. The attribute name gives the text MONDAY, and value gives the number 1 that was written after the equals sign.

Because the members are fixed, a value of type Day can only be one of the three members. A number such as 7 is not a Day, even though it could be stored in a variable. Plain enums are useful for that reason alone. The next sections add data to each member.

Giving each member data

Here is a coffee size that carries three facts: a short code, a volume in millilitres and a price in cents. Keeping the price in cents, which is a whole number, avoids the rounding errors that come with decimal fractions.

When a member's value is a tuple, Python passes its items to a method named __init__ that you write inside the enum. The method copies each item into an attribute of the member.

The output is:

CoffeeSize.SMALL SMALL S 240 $2.50
CoffeeSize.MEDIUM MEDIUM M 360 $3.00
CoffeeSize.LARGE LARGE L 480 $3.50
from_code l: $3.50
from_code XL: None

Read the file from top to bottom.

The line SMALL = ("S", 240, 250) creates the member SMALL, and its value is the tuple. Python calls __init__ once for each member, with the items of the tuple as arguments. The method copies each item into an attribute. So SMALL.code is "S", and LARGE.code is "L". Each member has its own copy of the three attributes.

The method price_text turns 350 cents into the text $3.50. The format :02d prints a whole number with at least two digits, so 5 cents becomes 05.

The class method from_code is a lookup. Given the text "l", it finds the member whose code matches, ignoring capital letters. If no member matches, it returns None. A caller who gets None has to check for it, and the missing case is easy to forget. The next section shows a lookup that makes the missing case harder to ignore.

The loop for size in cls works because an enum class can be looped over. It gives the members in the order they were written.

Lookups by name and by value

Python gives every enum two built-in lookups that you do not need to write.

• CoffeeSize["MEDIUM"] finds the member with that exact name. It is case-sensitive, so CoffeeSize["medium"] fails.
• CoffeeSize("M") finds the member whose value matches. Here the value is a tuple, so CoffeeSize("M") fails. Pass the whole tuple instead, as in CoffeeSize(("M", 360, 300)).

The next program uses name lookups, the match statement, and the checks Python makes when a lookup fails.

The output is:

Level.HIGH
Level.LOW
KeyError: 'MEDIUM'
ValueError: 9 is not a valid Level
high

The name lookup Level["HIGH"] raises KeyError for an unknown name. The value lookup Level(9) raises ValueError with a message that names the enum and the bad value. A program that reads a size from a user should catch these errors, or check the text first.

The match statement compares the value with each member. A case such as Level.LOW matches only that member, because a member is a single object. The equal sign is not used, and the dots name the member.

Members can be compared with is, and that is the normal comparison. They can also be compared with ==, which gives the same answer.

Methods on the enum

An enum can have as many methods as you need, and they can use the attributes of each member. The lookup from_code in the coffee example is a class method, so you call it on the class, as CoffeeSize.from_code("M").

Every enum member also has some attributes and methods from Python. These are the ones you will use most.

Looking up a member by name uses a dictionary that the enum keeps, so it stays fast for a long enum.

A lookup from a text code that is not the value needs a small method, like from_code. Python does not have a built-in lookup from a set of attributes. For a long list, build a dictionary once and read from it:

The output is:

CoffeeSize.MEDIUM
None

The dictionary comprehension builds a table from each code to its member, and get returns None for a missing key, the same as from_code. Building the table once is faster than looping over the members on every lookup, but for three members the loop is fine. Choose the table when the list is long, or when the same lookup happens many times.

Behaviour that differs by member

Attributes hold data. Sometimes each member needs its own behaviour, not just its own numbers. In Python, the easiest way is one method with a match statement, as the next example shows. Each case is the behaviour for one member.

The output is:

Operation.ADD + 6 and 3 -> 9
Operation.SUBTRACT - 6 and 3 -> 3
Operation.MULTIPLY * 6 and 3 -> 18
MULTIPLY None
Operation.SUBTRACT

The method apply checks which member it is running on, with match self. Each case returns the result for one operation. Adding a new member means adding one case. The match statement does not check that every member is handled, so a missing case would return None without any warning. A type checker can report it when the match is not complete.

The class method from_symbol looks up a member by its value, which here is the symbol. The value lookup Operation("-") gives the same member in one step, because the value is the symbol itself.

A function written in the class body becomes a method, not a member, so the match form above is the usual way to give members their own behaviour.

Ordering: IntEnum and plain enums

A plain enum member cannot be compared with <. The next example shows the error, and the fix, which is IntEnum, a version of enum whose members are also numbers.

The output is:

TypeError: '<' not supported between instances of 'Priority' and 'Priority'
True [<Rank.LOW: 1>, <Rank.HIGH: 2>]

A plain enum does not say which member comes first, so < is refused. An IntEnum member is a number, so it compares like one. Use IntEnum when the order of the members matters, such as a priority. Use a plain Enum when the values are only names.

If you need the position of a member in the list of members, use list(CoffeeSize).index(CoffeeSize.MEDIUM), which returns 1 for the second member. Do not save that position to a file, because reordering the members would change every saved number. Save the name instead.

Common mistakes

The following mistakes come up often. Each message below is the real output of Python.

Mistake 1: creating a member by calling the class with a name

A beginner may try to make a value in the normal way, by calling the class with a name. For an enum, the call looks up a value, and it does not create a new member.

The output is:

ValueError: 'LOW' is not a valid Level

The message says that "LOW" is not a value of Level. The members are the only values the enum has. Use Level.LOW to get a member, or Level["LOW"] to look one up by name.

Mistake 2: forgetting that a lookup by name can fail

Code that reads a size from a text file and looks it up by name will stop with KeyError on a typo such as "Medium " with a trailing space. Strip the text, and use the lookup that returns None for an unknown value, such as from_code, or catch KeyError where the text enters the program.

The output is:

KeyError: 'Medium '
CoffeeSize.MEDIUM

Mistake 3: two members with the same value

When two members have equal values, Python does not create a second member. The second name becomes an alias, which is another name for the first member. This is easy to miss, because the code still runs.

The output is:

True Shelf.KITCHEN
[<Shelf.KITCHEN: []>]

HALL is the same object as KITCHEN, and the list of members shows only one. Give each member a distinct value. Keep any list outside the enum, because every use of a member shares that one list.

Choosing between an enum, a dictionary and a class

Both enums and classes can hold a fixed set of named choices with data. They differ in how much the language does for you.

Here is the rule this lesson suggests.

Prefer an Enum for a fixed set of named choices, especially when each member carries data that belongs together, such as a code, a price and a volume. Keep __init__ simple and put the lookup in a class method. Build a dictionary when the list is long and the lookup is frequent.

Use an IntEnum when the members have a natural order. Use a plain list loaded at start-up when the choices can change without a new version of the program. The cost is that Python cannot check that you handled every choice, so your code must handle an unknown choice on its own.

Practice

Exercise 1. In the CoffeeSize enum above, add a new size EXTRA_LARGE with the code XL, 600 millilitres and a price of 425 cents. Then predict what CoffeeSize.from_code("XL").price_text() prints, and explain why no other method needs a change.

Show the answer

The change is one new line in the class body, after LARGE. The from_code method loops over all members, so it finds XL with no other change, and price_text() prints $4.25. The new member is found by its code, and every method that loops over the enum includes it automatically. A match statement that lists the members by name would need a new case, and a type checker would report a missing case if the match is complete.

Exercise 2. Write a class method largest() in CoffeeSize that returns the member with the most millilitres. Use a loop over the members. Do not use max or a sort.

Show the answer

The output is CoffeeSize.LARGE. The loop starts with no best member, and it replaces best each time it finds a member with more millilitres.

Exercise 3. Given an enum Operation with the members ADD, SUBTRACT and MULTIPLY, in that order, what does list(Operation).index(Operation.MULTIPLY) print? Answer with the number only, then say what the position counts from.

Show the answer

It prints 2. Positions count from 0, so ADD is 0, SUBTRACT is 1 and MULTIPLY is 2.

Exercise 4. Write an enum Shape with the members CIRCLE and SQUARE, and a method name_text() that returns "circle" or "square" by using a match statement on self. Print Shape.SQUARE.name_text().

Show the answer

The output is square.

Exercise 5. A teammate writes if shape == 1: for the Shape enum above. Explain what the comparison does and how to fix the line.

Show the answer

The comparison is always False, because a Shape member is not equal to the number 1. No error is raised, so the bug is silent. Change the line to if shape is Shape.CIRCLE:, which compares with the named member. If the number is really needed, read it with shape.value, so the conversion shows in the code.

Summary

• An enum is a fixed set of named members. Each member is one object that Python creates when the class is defined.
• A member can have attributes. Python passes the tuple value to __init__, which copies the values into attributes.
• Use Enum[name] to look up a member by its name, and Enum(value) to look up by its value. Both raise an error for an unknown key.
• Methods on an enum can use match self to give each member its own behaviour.
• Use IntEnum when the members need to be compared with < or sorted.
• Store only fixed values in members. Keep data that changes somewhere else.