Company logo

Quantlane tech blog

A 3x faster enum type for Python

Python's enum type is a useful building block for creating semantic types and constants in your programs. For example, instead of passing strings like 'BLUE' and 'GREEN' around, you can have Color.BLUE and Color.GREEN. The problem? It's kind of slow. We were using enums a lot in our code, until we noticed that enum overhead takes up single digit percentages of our CPU time! Luckily, it only took a few hours to write a much faster implementation with almost the same functionality. Enter fastenum.

Several years ago when fastenum was born it had a ~10x speed advantage over enum. Over the years the standard library implementation greatly improved, but fastenum remains up to 3.5x faster to this day.

fastenum will in many cases work as a drop-in replacement for enum after you install it with pip install fastenum:

import fastenum

class Color(fastenum.Enum):
    RED = 0
    BLUE = 1
    GREEN = 2

assert isinstance(Color.RED, Color)
assert Color.RED is Color['RED']
assert Color.BLUE != 1
assert Color.GREEN.value == 2

def is_red(c: Color) -> bool:
    return c is Color.RED

However, attribute access (every simple expression like Color.BLUE) is 3.5x faster in benchmarks. We think this is crucial since we do millions of attribute lookups on enums – it's pretty much impossible to use enums without attribute access. Here is a benchmark comparing fastenum to standard enum (lower times are better):

------------------------- benchmark 'test_attribute_access': 2 tests -------------------------
Name (time in ns)                       Mean              StdDev              Median
----------------------------------------------------------------------------------------------
test_attribute_access[fastenum]     127.5973 (1.0)       62.0568 (1.0)      116.2900 (1.0)
test_attribute_access[enum]         447.2529 (3.51)     885.4300 (14.27)    408.0000 (3.51)
----------------------------------------------------------------------------------------------

Other common operations are 1.5–2.5x faster, too: this is the 'call' access (Color(2)), 'getitem' access (Color['BLUE']), and enum iteration (for c in Color:):

--------------------------- benchmark 'test_call': 2 tests ---------------------------
Name (time in ns)             Mean                StdDev              Median
--------------------------------------------------------------------------------------
test_call[fastenum]       425.0350 (1.0)        145.8800 (1.0)      401.7500 (1.0)
test_call[enum]         1,066.1683 (2.51)     1,076.5092 (7.38)     957.0000 (2.38)
--------------------------------------------------------------------------------------

------------------------- benchmark 'test_getitem_access': 2 tests -------------------------
Name (time in ns)                     Mean              StdDev              Median
--------------------------------------------------------------------------------------------
test_getitem_access[fastenum]     241.7980 (1.0)       99.1632 (1.0)      227.1364 (1.0)
test_getitem_access[enum]         606.9534 (2.51)     730.2464 (7.36)     562.0000 (2.47)
--------------------------------------------------------------------------------------------

---------------------- benchmark 'test_iter': 2 tests ----------------------
Name (time in us)         Mean            StdDev            Median
----------------------------------------------------------------------------
test_iter[fastenum]     1.5749 (1.0)      1.3203 (1.0)      1.4580 (1.0)
test_iter[enum]         2.3743 (1.51)     1.5423 (1.17)     2.2030 (1.51)
----------------------------------------------------------------------------

We also cache the enum members' __hash__ and __repr__ (easy to do since enum members are constant). The __hash__ cache is useful in particular, since we tend to use enums as dictionary keys in some cases.

We also developed a mypy plugin for fastenum which you can enable in mypy.ini:

[mypy]
plugins = fastenum.mypy_plugin:plugin

This is necessary for mypy to understand fastenum the same way it understands enum (for which it has built-in support). Note you may experience odd mypy crashes relating to the mypy cache when using this plugin. It's an unfortunate bug we have not yet tracked down. If you run into it, it is unfortunately necessary to disable the mypy cache by setting cache_dir = /dev/null 🤦‍♂️

How we did it & what the trade-offs are

fastenum came into being as a simple experiment where I took the original enum implementation and started deleting features and simplifying and caching what was left. This worked surprisingly well – we got a 10x speedup out of the gate. (This advantage gradually diminished to 3.5x as the standard enum performance got better.) My awesome colleagues then kept improving fastenum over the years and, most importantly, made it compatible with mypy (we love mypy at Quantlane).

fastenum is a stripped-down racecar next to your comfortable enum sedan. There is no support for automatic values, unique value checks, aliases, custom __init__ implementations on members, IntEnum, Flag, or the functional API. If you require any of these features it's probably best to just use enum. That said, we are happy to accept pull requests with such features, provided they do not impact base fastenum performance!

Where to get it

fastenum is open-sourced under Apache Licence 2.0 and is tested to work with Python 3.7 and 3.8. It is available on PyPI with a simple pip install fastenum and the source code is on our GitLab.

To run the benchmarks yourself, install dev dependencies (pip install -r dev-requirements.txt) and run pytest: PYTHONPATH=. pytest.

Go give it a try and let us know how it works for you! 💌 code@quantlane.com

Quantlane Written by Vita Smid on May 28, 2021. We have more articles like this on our blog. If you want to be notified when we publish something about Python or finance, sign up for our newsletter:
Vita Smid

Hi, my name is Vita. I co-founded Quantlane together with a few stock traders in 2014. I developed the first version of our Python trading platform from the ground up. After our engineering team started to grow I became Quantlane's CTO. This was quite new for me, as prior to this I had been working as a freelance developer and had a degree in financial mathematics. In 2022, I left the company to seek new adventures.