Using enumerated types in Python
By John Lekberg on June 06, 2020.
This week's post is about using enumerated types (enums) in Python. You will learn:
- How to use enums to manage data that takes on a finite set of states.
- How to impose a custom ordering on data.
- How to represent a combination of states using
Flag
andIntFlag
.
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:
- The data that I'm working with has a small number of states (e.g. "pass" or "fail").
- The data has an ordering that doesn't work with strings (e.g. "S" before "A" in a tier list).
- The data can be represented as a set of boolean flags (e.g. the flag register of a Zilog Z80 microprocessor).
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:
Enum
IntEnum
IntFlag
Flag
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.)