ASPN : Python Cookbook : Enums for Python

I once tried to give Python something like C's enums, as described
here: http://groups.google.com/groups?selm=G6qzLy.6Fo%40world.std.com
That approach tried to assign to a dictionary returned by the locals()
function, intending that such assignments would become class attributes.
The Tim-bot explained to me the errors of my ways. The quest for the
perfect Python enum goes on.

Source: Text Source

import types, string, pprint, exceptions

class EnumException(exceptions.Exception):
    pass

class Enumeration:
    def __init__(self, name, enumList):
        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 x in uniqueNames:
                raise EnumException, "enum name is not unique: " + x
            if i in uniqueValues:
                raise EnumException, "enum value is not unique for " + x
            uniqueNames.append(x)
            uniqueValues.append(i)
            lookup[x] = i
            reverseLookup[i] = x
            i = i + 1
        self.lookup = lookup
        self.reverseLookup = reverseLookup
    def __getattr__(self, attr):
        if not self.lookup.has_key(attr):
            raise AttributeError
        return self.lookup[attr]
    def whatis(self, value):
        return self.reverseLookup[value]

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)

# Notice that car's and bug's attributes don't include any of the
# enum machinery, because that machinery is all CLASS attributes and
# not INSTANCE attributes. So you can generate thousands of cars and
# bugs with reckless abandon, never worrying that time or memory will
# be wasted on redundant copies of the enum stuff.

print car.__dict__
print bug.__dict__
pprint.pprint(Volkswagen.__dict__)
pprint.pprint(Insect.__dict__)
""")

Discussion:

In C, enums allow you to declare a bunch of constants with unique values,
without necessarily specifying the actual values (except in cases where you
need to). Python has an accepted idiom that's fine for very small numbers of
constants (A, B, C, D = range(4)) but it doesn't scale well to large numbers,
and it doesn't allow you to specify values for some constants while leaving
others unspecified. This approach does those things, while verifying that all
values (specified and unspecified) are unique. Enum values then are attributes
of an Enumeration class (Volkswagen.BEETLE, Volkswagen.PASSAT, etc.).

Add comment

Number of comments: 3

A different approach, Michael Radziej, 2001/08/29
I like it more complicated :-) The following class allows you to use enums as like colors.red, to convert like colors["red"], get the string value like someColor.asString and some other nifty things. Of course, the overhead is higher. Here it is ...

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

class Enum emulates Enumerations.

Enum instances contain _EnumNodes. These have two attributes:
asString and asInt.

Create one with

    Enum(list,*startvalue)
or: Enum(string,*startvalue)       (default for startvalue is 0)

e.g.:

WD = Enum(["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"], 1)
WD = Enum["Monday Tuesday Wednesday Thursday Friday Saturday Sunday", 1]

Typical use:

workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                             # for Enums, hence "irange"
                                             # equivalent: workdays = WD[1:6]
for i in WD.each():
    if i in workdays:
        print i.asInt, i.asString + " is a work day"
    else:
        print i.asInt, i.asString + " is a weekend day."
print "There are ", len(WD), "days per week"
if "Monday" in WD: print "Monday is a valid name"
if not "August" in WD: print "August is not"

some systematic examples:

WD.Monday.asString       --> 'Monday'
WD.Monday.asInt          --> 1
WD.stringToInt('Monday') --> 1
WD[2].asString           --> 'Tuesday'
WD.intToString(2)        --> 'Tuesday'
WD["Tuesday"].asInt      --> 2
WD.Saturday > WD.Tuesday --> 1
"Monday" in WD --> 1
6 in WD --> 1
0 in WD --> 0
WD.each() gives you a list of all EnumNodes.
WD.eachString() --> find out by yourself!

Note: You cannot create EnumNodes other than creating an Enum.
      The EnumNode class is considered private.
      Never not try to change attributes of these objects.

"""

################################################################

import types

class _EnumNode:
    def __init__(self,i,name):
        self.asInt=i
        self.asString=name

    def __cmp__(left,right):
        return cmp(left.asInt, right.asInt)

    def __str__(self):
        return self.asString

    def __repr__(self):
        return "("+str(self.asInt)+":"+self.asString+")"

    def __hash__(self):
        return self.asInt

class Enum:
    def __init__(self,stringList,start=0):
        if type(stringList)==types.StringType:
            stringList = stringList.split()
        self._start = start
        self._byString = {}
        self._byInt = [ None ] * (start + len(stringList))
        for i in range(len(stringList)):
            node = _EnumNode(i+start,stringList[i])
            self._byInt[i+start] = node
            self._byString[node.asString] = node

    def addAlternate(self,node,aString):
        self._byString[aString] = node

    def intToString(self,num):
        return self._byInt[num].asString

    def stringToInt(self,name):
        node = self._byString.get(name,None)
        if node: return node.asInt
        else: return None

    def __contains__(self,key):
        if type(key)==types.IntType:
            return key >= self._start and key I like it more complicated :-)
The following class allows you to use enums as like
colors.red, to convert like colors["red"], get the string value
like someColor.asString and some other
nifty things. Of course, the overhead is higher.

Here it is ...

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

class Enum emulates Enumerations.

Enum instances contain _EnumNodes. These have two attributes:
asString and asInt.

Create one with

    Enum(list,*startvalue)
or: Enum(string,*startvalue)       (default for startvalue is 0)

e.g.:

WD = Enum(["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"], 1)
WD = Enum["Monday Tuesday Wednesday Thursday Friday Saturday Sunday", 1]

Typical use:

workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                             # for Enums, hence "irange"
                                             # equivalent: workdays = WD[1:6]
for i in WD.each():
    if i in workdays:
        print i.asInt, i.asString + " is a work day"
    else:
        print i.asInt, i.asString + " is a weekend day."
print "There are ", len(WD), "days per week"
if "Monday" in WD: print "Monday is a valid name"
if not "August" in WD: print "August is not"

some systematic examples:

WD.Monday.asString       --> 'Monday'
WD.Monday.asInt          --> 1
WD.stringToInt('Monday') --> 1
WD[2].asString           --> 'Tuesday'
WD.intToString(2)        --> 'Tuesday'
WD["Tuesday"].asInt      --> 2
WD.Saturday > WD.Tuesday --> 1
"Monday" in WD --> 1
6 in WD --> 1
0 in WD --> 0
WD.each() gives you a list of all EnumNodes.
WD.eachString() --> find out by yourself!

Note: You cannot create EnumNodes other than creating an Enum.
      The EnumNode class is considered private.
      Never not try to change attributes of these objects.

"""

################################################################

import types

class _EnumNode:
    def __init__(self,i,name):
        self.asInt=i
        self.asString=name

    def __cmp__(left,right):
        return cmp(left.asInt, right.asInt)

    def __str__(self):
        return self.asString

    def __repr__(self):
        return "("+str(self.asInt)+":"+self.asString+")"

    def __hash__(self):
        return self.asInt

class Enum:
    def __init__(self,stringList,start=0):
        if type(stringList)==types.StringType:
            stringList = stringList.split()
        self._start = start
        self._byString = {}
        self._byInt = [ None ] * (start + len(stringList))
        for i in range(len(stringList)):
            node = _EnumNode(i+start,stringList[i])
            self._byInt[i+start] = node
            self._byString[node.asString] = node

    def addAlternate(self,node,aString):
        self._byString[aString] = node

    def intToString(self,num):
        return self._byInt[num].asString

    def stringToInt(self,name):
        node = self._byString.get(name,None)
        if node: return node.asInt
        else: return None

    def __contains__(self,key):
        if type(key)==types.IntType:
            return key >= self._start and key

Add comment

Re: A different approach, Martin Miller,Martin Miller, 2005/03/21
On 2001/08/29 Michael Radziej wrote:
> I like it more complicated :-)

That's fine, as long as the code isn't broken and is complete, which does not seem to be the case in what was posted. Specifically there was a crucial line missing from Enum.__init__(), the cleverness in the if in Enum.__contains__() messed up when the default start values was used, and a number other methods need for the examples were missing.

Below is a complete and functional version that was tested with Python 2.4. There are comments for most of the changes made.

Additional observations: The list contained in the _byInt attribute could be quite long if the starting integer value is is a big number. Also, some of the methods are fairly inefficient in those circumstances.

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
"""
class Enum emulates Enumerations.

Enum instances contain _EnumNodes. These have two attributes:
asString and asInt.

Create one with

    Enum(list,*startvalue)
or: Enum(string,*startvalue)       (default for startvalue is 0)

e.g.:

WD = Enum(["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"], 1)
WD = Enum["Monday Tuesday Wednesday Thursday Friday Saturday Sunday", 1]

Typical use:

workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                             # for Enums, hence "irange"
                                             # equivalent: workdays = WD[1:6]
for i in WD.each():
    if i in workdays:
        print i.asInt, i.asString + " is a work day"
    else:
        print i.asInt, i.asString + " is a weekend day."
print "There are ", len(WD), "days per week"
if "Monday" in WD: print "Monday is a valid name"
if not "August" in WD: print "August is not"

some systematic examples:

WD.Monday.asString       --> 'Monday'
WD.Monday.asInt          --> 1
WD.stringToInt('Monday') --> 1
WD[2].asString           --> 'Tuesday'
WD.intToString(2)        --> 'Tuesday'
WD["Tuesday"].asInt      --> 2
WD.Saturday > WD.Tuesday --> 1
"Monday" in WD --> 1
6 in WD --> 1
0 in WD --> 0
WD.each() gives you a list of all EnumNodes.
WD.eachString() --> find out by yourself!

Note: You cannot create EnumNodes other than creating an Enum.
      The EnumNode class is considered private.
      Never not try to change attributes of these objects.

"""

################################################################

import types

class _EnumNode:
    def __init__(self,i,name):
        self.asInt=i
        self.asString=name

    def __cmp__(left,right):
        return cmp(left.asInt, right.asInt)

    def __str__(self):
        return self.asString

    def __repr__(self):
        return "("+str(self.asInt)+":"+self.asString+")"

    def __hash__(self):
        return self.asInt

class Enum:
    def __init__(self,stringList,start=0):
        if type(stringList)==types.StringType:
            stringList = stringList.split()
        self._start = start
        self._byString = {}
        self._byInt = [ None ] * (start + len(stringList))
        for i in range(len(stringList)):
            node = _EnumNode(i+start,stringList[i])
            setattr(self, stringList[i], node)
            self._byInt[i+start] = node
            self._byString[node.asString] = node

    def addAlternate(self,node,aString): # note: doesn't update '_byInt'
        self._byString[aString] = node

    def intToString(self,num):
        return self._byInt[num].asString

    def stringToInt(self,name):
        node = self._byString.get(name,None)
        if node: return node.asInt
        else: return None

    def __contains__(self,key):
        if type(key)==types.IntType:
            return key >= self._start # removed 'and key' (didn't allow a 0 start value)
        else: # else needed for strings
            return key in self._byString

    # missing methods...

    def __len__(self):
        return len(self._byInt)-self._start

    def __getitem__(self,key):
        if type(key)==types.StringType:
            return self._byString[key]
        else:
            return self._byInt[key]

    def irange(self, begin, end):
        return [self._byInt[i] for i in range(begin.asInt, end.asInt+1)]

    def each(self):
        return [node for node in self._byInt if node is not None]

    def eachString(self):
        return [node.asString for node in self.each()]


if __name__ == '__main__':
    # Typical use:

    WD = Enum("Monday Tuesday Wednesday Thursday Friday Saturday Sunday", 1)

    workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                                 # for Enums, hence "irange"
                                                 # equivalent: workdays = WD[1:6]
    print "workdays: ", workdays
    print "WD[1:6]: ", WD[1:6]

    for i in WD.each():
        if i in workdays:
            print i.asInt, i.asString + " is a work day"
        else:
            print i.asInt, i.asString + " is a weekend day."

    print "There are", len(WD), "days per week"
    if "Monday" in WD: print "Monday is a valid name"
    if not "August" in WD: print "August is not"

    # some systematic examples:

    print "WD.Monday.asString:", WD.Monday.asString
    print "WD.Monday.asInt:", WD.Monday.asInt
    print "WD.stringToInt('Monday'):", WD.stringToInt('Monday')
    print "WD[2].asString:", WD[2].asString
    print "WD.intToString(2):", WD.intToString(2)
    print "WD['Tuesday'].asInt:", WD['Tuesday'].asInt
    print "WD.Saturday > WD.Tuesday:", WD.Saturday > WD.Tuesday
    print "'Monday' in WD:", "Monday" in WD
    print "6 in WD:", 6 in WD
    print "0 in WD:", 0 in WD
    print "WD.each():", WD.each()
    print "WD.eachString():", WD.eachString()

Add comment

Specified values should be considered before unspecified, Will Ware, 2007/02/21
This fixes a minor bug that would have become apparent if my test case had used specified values that might have collided with the normal numbering (e.g. 4 for Volkswagen.THING instead of 400).

class Enumeration:
    def __init__(self, name, enumList):
        self.__doc__ = name
        lookup = { }
        reverseLookup = { }
        uniqueNames = [ ]
        self._uniqueValues = uniqueValues = [ ]
        self._uniqueId = 0
        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 x in uniqueNames:
                    raise EnumException, "enum name is not unique: " + x
                if i in uniqueValues:
                    raise EnumException, "enum value is not unique for " + x
                uniqueNames.append(x)
                uniqueValues.append(i)
                lookup[x] = i
                reverseLookup[i] = x
        for x in enumList:
            if type(x) != types.TupleType:
                if type(x) != types.StringType:
                    raise EnumException, "enum name is not a string: " + x
                if x in uniqueNames:
                    raise EnumException, "enum name is not unique: " + x
                uniqueNames.append(x)
                i = self.generateUniqueId()
                uniqueValues.append(i)
                lookup[x] = i
                reverseLookup[i] = x
        self.lookup = lookup
        self.reverseLookup = reverseLookup
    def generateUniqueId(self):
        while self._uniqueId in self._uniqueValues:
            self._uniqueId += 1
        n = self._uniqueId
        self._uniqueId += 1
        return n
    def __getattr__(self, attr):
        if not self.lookup.has_key(attr):
            raise AttributeError
        return self.lookup[attr]
    def whatis(self, value):
        return self.reverseLookup[value]

Add comment

Highest rated recipes:

1. A simple XML-RPC server

2. Web service accessible ...

3. IPy Notify

4. Changing return value ...

5. Quantum Superposition

6. Pickle objects under ...

7. Generalized delegates ...

8. Reorder a sequence (uses ...

9. Setting Win32 System ...

10. ObjectMerger