Return to Blog

Using enumerated types in Python

By John Lekberg on June 06, 2020.


Hacker News discussion


This week's post is about using enumerated types (enums) in Python. You will learn:

What are enums?

Enumerated types (enums) are sets of unique values. E.g. A tier list, with S-tier being the best, followed by A, B, C, D, and F:

import enum

Tier = enum.IntEnum("Tier", ["S", "A", "B", "C", "D", "F"])
list(Tier)
[<Tier.S: 1>, <Tier.A: 2>, <Tier.B: 3>, <Tier.C: 4>, <Tier.D: 5>,
 <Tier.F: 6>]
Tier.S < Tier.C
True
sorted([Tier.B, Tier.D, Tier.S, Tier.A, Tier.F, Tier.C])
[<Tier.S: 1>, <Tier.A: 2>, <Tier.B: 3>, <Tier.C: 4>, <Tier.D: 5>,
 <Tier.F: 6>]

In this case, using an IntEnum object is easier than using a string because the tiers compare correctly:

Tier.S < Tier.C
True
"S" < "C"
False
sorted([Tier.B, Tier.D, Tier.S, Tier.A, Tier.F, Tier.C])
[<Tier.S: 1>, <Tier.A: 2>, <Tier.B: 3>, <Tier.C: 4>, <Tier.D: 5>,
 <Tier.F: 6>]
sorted(["B", "D", "S", "A", "F", "C"])
['A', 'B', 'C', 'D', 'F', 'S']

And spelling mistakes are caught:

Tier.Z < Tier.Q
AttributeError: Z
"Z" < "Q"
False

Using an IntEnum object is easier than using an int because the IntEnum allows you to write mnemonics for the values:

Tier.S < Tier.C
True
1 < 4
True
sorted([Tier.B, Tier.D, Tier.S, Tier.A, Tier.F, Tier.C])
[<Tier.S: 1>, <Tier.A: 2>, <Tier.B: 3>, <Tier.C: 4>, <Tier.D: 5>,
 <Tier.F: 6>]
sorted([3, 5, 1, 2, 6, 4])
[1, 2, 3, 4, 5, 6]

When should I use enums?

I use enums when:

Here's code for each of these examples:


import random
import collections

Test = enum.Enum("Test", ["Pass", "Fail"])

row = lambda: [
    random.choice(list(Test)),
    random.choice(["Bobby", "Linda", "John"])
]
data = [ row() for _ in range(10) ]
data
[[<Test.Pass: 1>, 'John'],
 [<Test.Pass: 1>, 'Bobby'],
 [<Test.Fail: 2>, 'Linda'],
 [<Test.Fail: 2>, 'Bobby'],
 [<Test.Pass: 1>, 'Linda'],
 [<Test.Pass: 1>, 'Bobby'],
 [<Test.Pass: 1>, 'Linda'],
 [<Test.Pass: 1>, 'Bobby'],
 [<Test.Pass: 1>, 'Bobby'],
 [<Test.Pass: 1>, 'Bobby']]
[
    student
    for test, student in data
    if test == Test.Fail
]
['Linda', 'Bobby']

import random
import collections

Tier = enum.IntEnum("Tier", ["S", "A", "B", "C", "D", "F"])

data_tiers = random.choices(list(Tier), k=9)
data_names = ["Mario", "Link", "Samus", "Yoshi", "Kirby",
              "Fox", "Pikachu", "Ness", "Luigi"]
data = list(zip(data_tiers, data_names))
data
[(<Tier.B: 3>, 'Mario'),
 (<Tier.A: 2>, 'Link'),
 (<Tier.F: 6>, 'Samus'),
 (<Tier.B: 3>, 'Yoshi'),
 (<Tier.C: 4>, 'Kirby'),
 (<Tier.D: 5>, 'Fox'),
 (<Tier.A: 2>, 'Pikachu'),
 (<Tier.S: 1>, 'Ness'),
 (<Tier.A: 2>, 'Luigi')]
rankings = collections.defaultdict(set)

for tier, name in data:
    rankings[tier].add(name)

for tier in Tier:
    names = sorted(rankings[tier])
    print(f"{tier.name}-tier: {', '.join(names)}")
S-tier: Ness
A-tier: Link, Luigi, Pikachu
B-tier: Mario, Yoshi
C-tier: Kirby
D-tier: Fox
F-tier: Samus

import random

state = 0
for flag in Z80_Flag:
    # 50% chance that p is True.
    p = random.random() < 0.5
    if p:
        state |= flag

state
<Z80_Flag.Sign|Parity_Overflow|Subtract: 134>
for flag in Z80_Flag:
    if state & flag:
        status = "#"
    else:
        status = "."
    print(f"{flag.name:16} {status}")
Carry            .
Subtract         #
Parity_Overflow  #
Half_Carry       .
Zero             .
Sign             #

Python's support for enums

Python supports four different types of enums via the enum module:

What are the differences between these four?

Enum and Flag do not mix with integers. IntEnum and IntFlag do mix with integers.

import enum

Test = enum.Enum("Test", ["Pass", "Fail"])
IntTest = enum.IntEnum("IntTest", ["Pass", "Fail"])

Test.Pass + 0
TypeError: unsupported operand type(s) for +: 'Test' and 'int'
IntTest.Pass + 0
1
Z80 = enum.Flag("Z80", ["Carry", "Subtract", "Parity", "Zero"])
IntZ80 = enum.IntFlag("IntZ80", ["Carry", "Subtract", "Parity",
                                 "Zero"])

Z80.Carry | 2
TypeError: unsupported operand type(s) for |: 'Z80' and 'int'
IntZ80.Carry | 2
<IntZ80.Subtract|Carry: 3>

Enum and IntEnum represent states that can't be combined. Flag and IntFlag represent states that can be combined.

Test = enum.Enum("Test", ["Pass", "Fail"])
Z80 = enum.Flag("Z80", ["Carry", "Subtract", "Parity", "Zero"])

Test.Pass | Test.Fail
TypeError: unsupported operand type(s) for |: 'Test' and 'Test'
Z80.Carry | Z80.Zero
<Z80.Zero|Carry: 9>

In conclusion...

In this week's post you learned how to use enums to manage data that takes on a finite set of states. Enums allow you to impose a custom ordering on data and represent a combination of states using Flag and IntFlag.

My challenge to you:

Create an IntFlag that represents the input status of a standard NES controller.

If you enjoyed this week's post, share it with your friends and stay tuned for next week's post. See you then!


(If you spot any errors or typos on this post, contact me via my contact page.)