Python Cookbook

Simulating Enumerations in Python

Credit: Will Ware

Problem

You want to define an enumeration in the spirit of C’s enum type.

Solution

Python’s introspection facilities let you add a version of enum, even though Python, as a language, does not support this construct:

import types, string, pprint, exceptions

class EnumException(exceptions.Exception):
    pass

class Enumeration:
    def _ _init_ _(self, name, enumList, valuesAreUnique=1):
        self._ _doc_ _ = name
        lookup = { }
        reverseLookup = { }
        i = 0
        uniqueNames = {}
        uniqueValues = {}
        for x in enumList:
            if type(x) == types.TupleType:
                x, i = x
            if type(x) != types.StringType:
                raise EnumException, "enum name is not a string: " + x
            if type(i) != types.IntType:
                raise EnumException, "enum value is not an integer: " + i
            if uniqueNames.has_key(x):
                raise EnumException, "enum name is not unique: " + x
            if valuesAreUnique and uniqueValues.has_key(i):
                raise EnumException, "enum value is not unique for " + x
            uniqueNames[x] = 1
            uniqueValues[i] = 1
            lookup[x] = i
            reverseLookup[i] = x
            i = i + 1
        self.lookup = lookup
        self.reverseLookup = reverseLookup
    def _ _getattr_ _(self, attr):
        try: return self.lookup[attr]
        except KeyError: raise AttributeError
    def whatis(self, value):
        return self.reverseLookup[value]

Discussion

In C, enum lets you declare several constants, typically with unique values (although you can also explicitly arrange for a value to be duplicated under two different names), without necessarily specifying the actual values (except when you want it to).

Python has an accepted idiom that’s fine for small numbers of constants:

A, B, C, D = range(4)

But this idiom doesn’t scale well to large numbers and doesn’t allow you to specify values for some constants while leaving others to be determined automatically. This recipe provides for all these niceties, while optionally verifying that all values (specified and unspecified) are unique. Enum values are attributes of an Enumeration class (Volkswagen.BEETLE, Volkswagen.PASSAT, etc.). A further feature, missing in C but really quite useful, is the ability to go from the value to the corresponding name inside the enumeration (of course, the name you get is somewhat arbitrary for those enumerations in which you don’t constrain values to be unique).

This recipe’s Enumeration class has an instance constructor that accepts a string argument to specify the enumeration’s name and a list argument to specify the names of all values for the enumeration. Each item of the list argument can be a string (to specify that the value named is one more than the last value used), or else a tuple with two items (the string that is the value’s name and the value itself, which must be an integer). The code in this recipe relies heavily on strict type-checking to find out which case applies, but the recipe’s essence would not change by much if the checking was performed in a more lenient way (e.g., with the isinstance built-in function).

Therefore, each instance is equipped with two dictionaries: self.lookup to map names to values and self.reverselookup to map values back to the corresponding names. The special method _ _getattr_ _ lets names be used with attribute syntax (e.x is mapped to e.lookup['x']), and the whatis method allows reverse lookups (i.e., finds a name, given a value) with comparable syntactic ease.

Here’s an example of how you can use this Enumeration class:

if _ _name_ _ == '_ _main_ _':

    Volkswagen = Enumeration("Volkswagen",
        ["JETTA", "RABBIT", "BEETLE", ("THING", 400), "PASSAT", "GOLF",
         ("CABRIO", 700), "EURO_VAN", "CLASSIC_BEETLE", "CLASSIC_VAN"
         ])

    Insect = Enumeration("Insect",
        ["ANT", "APHID", "BEE", "BEETLE", "BUTTERFLY", "MOTH", "HOUSEFLY",
         "WASP", "CICADA", "GRASSHOPPER", "COCKROACH", "DRAGONFLY"
         ])

    def demo(lines):
        previousLineEmpty = 0
        for x in string.split(lines, "\n"):
            if x:
                if x[0] != '#':
                    print ">>>", x; exec x; print
                    previousLineEmpty = 1
                else:
                    print x
                    previousLineEmpty = 0
            elif not previousLineEmpty:
                print x
                previousLineEmpty = 1

    def whatkind(value, enum):
        return enum._ _doc_ _ + "." + enum.whatis(value)

    class ThingWithType:
        def _ _init_ _(self, type):
            self.type = type

    demo("""
    car = ThingWithType(Volkswagen.BEETLE)
    print whatkind(car.type, Volkswagen)
    bug = ThingWithType(Insect.BEETLE)
    print whatkind(bug.type, Insect)
    print car._ _dict_ _
    print bug._ _dict_ _
    pprint.pprint(Volkswagen._ _dict_ _)
    pprint.pprint(Insect._ _dict_ _)
    """)

Note that attributes of car and bug don’t include any of the enum machinery, because that machinery is held as class attributes, not as instance attributes. This means you can generate thousands of car and bug objects with reckless abandon, never worrying about wasting time or memory on redundant copies of the enum stuff.

See Also

Recipe 5.16, which shows how to define constants in Python; documentation on _ _getattr_ _ in the Language Reference.