ENOSUCHBLOG

Programming, philosophy, pedaling.


Totally ordered enums in Python with ordered_enum

Mar 2, 2020

Tags: programming, devblog

Yet another library announcement, and a Python one at that: ordered_enum.

ordered_enum solves a singular problem: providing a totally-ordered enumeration type in Python. Piggybacking on Python’s lovely enum.Enum, it provides two independent ordering solutions: one based on definition order (OrderedEnum), and another based on value order (ValueOrderedEnum).

Why?

Because Python’s enum.Enum does not provide ordering by default, ostensibly to eliminate C-style (mis-)treatment of enumerations as thin wrappers over integral types. ordered_enum re-adds ordering to Python’s enums without exposing value restrictions.

More concretely: enums are a nice way to encapsulate abstract state. Ordering supplies a notion of “forwardness” or “weight” in a set of enumerated states. This has a number of useful applications, like sorting complex objects by an enumerated property:

from ordered_enum import OrderedEnum


class Shape(OrderedEnum):
    Circle = 1
    Square = 2
    Triangle = 3


class Foo:
    def __init__(self, shape):
        self.shape = shape


foos = [
    Foo(Shape.Triangle), Foo(Shape.Circle), Foo(Shape.Triangle), Foo(Shape.Square)
]
sorted(foos, key=lambda foo: foo.shape)

or propagating a state machine via some notion of “forwardness”:

from ordered_enum import OrderedEnum


class State(OrderedEnum):
    Created: int = 0
    Waiting: int = 1
    Running: int = 2
    Done: int = 3


def tick():
    state = State.Created

    while tick():
        # Our enum's ordering allows us to use this as a shorthand for
        # `state == State.Created or state == State.Waiting or state == State.running`
        if state < State.Done:
            pass

The API

ordered_enum provides two totally-ordered enum classes, both of which inherit from enum.Enum:

Observe that OrderedEnum is the more permissive of the two: because it uses only the definition order, the values defined in an OrderedEnum can be of heterogeneous types. That means that this, even if ill-advised, works perfectly well:

from ordered_enum import OrderedEnum


class Animal(OrderedEnum):
    Cat = "cat"
    Dog = 2
    Iguana = {"dictionaries": "can be values too"}
    Human = {"and", "so", "can", "sets"}


Animal.Human > Animal.Cat  # True
Animal.Iguana > Animal.Human  # False

By contrast, ValueOrderedEnum assumes the uniqueness (and comparability) of its values. The former can be strongly enforced with the @enum.unique decorator, but is not by default:

import enum
from ordered_enum import ValueOrderedEnum


@enum.unique
class Humor(ValueOrderedEnum):
    Blood = 4
    Phlegm = 3
    YellowBile = 2
    BlackBile = 1


Humor.Blood > Humor.Phlegm  # True
Humor.BlackBile > Humor.YellowBile  # False

How it works

Both OrderedEnum and ValueOrderedEnum essentially boil down to a single custom method: __lt__. Combined with functools.total_ordering and the fact that enum.Enum provides __eq__, each has a straightforward implementation.

For OrderedEnum:

@functools.total_ordering
class OrderedEnum(Enum):
    @classmethod
    @functools.lru_cache(None)
    def _member_list(cls):
        return list(cls)

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            member_list = self.__class__._member_list()
            return member_list.index(self) < member_list.index(other)
        return NotImplemented

…and for ValueOrderedEnum:

@functools.total_ordering
class ValueOrderedEnum(Enum):
    def __lt__(self, other):
        if self.__class__ is other.__class__:
            return self.value < other.value
        return NotImplemented

Installation

ordered_enum is available on PyPI and installable via pip. It supports Python 3.6 and newer:

$ pip3 install ordered_enum