ActiveState Powered by ActiveState

Recipe 67107: 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.

Python
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
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.).

Comments

  1. 1. At 8 a.m. on 29 aug 2001, Michael Radziej said:

    A different approach. 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
    

    (comment continued...)

  2. 2. At 8 a.m. on 29 aug 2001, Michael Radziej said:

    (...continued from previous comment)

    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 < len(self._byInt)
            else:
                return self._byString.has_key(key)
    
        def __getattr__(self,name):
            return self._byString[name]
    
        def __len__(self):
            return len(self._byInt) - self._start
    
        def __getitem__(self,key):
            if type(key)==types.IntType:
                return self._byInt[key]
            elif type(key)==types.SliceType:
                return self._byInt[max(self._start,key.start):key.stop]
            else:
                return self._byString[key]
    
        def __repr__(self):
            return str(self._byInt[self._start:])
    
        def irange(self, start=None, stop=None, step=1):
            if start:
                startInt = start.asInt
            else:
                startInt = self._start
            if stop:
                stopInt = stop.asInt + 1
            else:
                stopInt = len(self._byInt)
            return [self._byInt[i] for i in range(startInt, stopInt,step)]
    
        def each(self):
            return self._byInt[self._start:]
    
        def eachString(self):
            return [n.asString for n in self._byInt[self._start:]]
    
  3. 3. At 1:58 p.m. on 21 mar 2005, Martin Miller said:

    Re: A different approach. 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.
    &quot;&quot;&quot;
    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([&quot;Monday&quot;, &quot;Tuesday&quot;, &quot;Wednesday&quot;, &quot;Thursday&quot;, &quot;Friday&quot;, &quot;Saturday&quot;, &quot;Sunday&quot;], 1)
    WD = Enum[&quot;Monday Tuesday Wednesday Thursday Friday Saturday Sunday&quot;, 1]
    
    Typical use:
    
    workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                                 # for Enums, hence &quot;irange&quot;
                                                 # equivalent: workdays = WD[1:6]
    for i in WD.each():
        if i in workdays:
            print i.asInt, i.asString + &quot; is a work day&quot;
        else:
            print i.asInt, i.asString + &quot; is a weekend day.&quot;
    print &quot;There are &quot;, len(WD), &quot;days per week&quot;
    if &quot;Monday&quot; in WD: print &quot;Monday is a valid name&quot;
    if not &quot;August&quot; in WD: print &quot;August is not&quot;
    
    some systematic examples:
    
    WD.Monday.asString       --&gt; 'Monday'
    WD.Monday.asInt          --&gt; 1
    WD.stringToInt('Monday') --&gt; 1
    WD[2].asString           --&gt; 'Tuesday'
    WD.intToString(2)        --&gt; 'Tuesday'
    WD[&quot;Tuesday&quot;].asInt      --&gt; 2
    WD.Saturday &gt; WD.Tuesday --&gt; 1
    &quot;Monday&quot; in WD --&gt; 1
    6 in WD --&gt; 1
    0 in WD --&gt; 0
    WD.each() gives you a list of all EnumNodes.
    WD.eachString() --&gt; 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.
    

    (comment continued...)

  4. 4. At 1:58 p.m. on 21 mar 2005, Martin Miller said:

    (...continued from previous comment)

    &quot;&quot;&quot;
    
    ################################################################
    
    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 &quot;(&quot;+str(self.asInt)+&quot;:&quot;+self.asString+&quot;)&quot;
    
        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 &gt;= 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(&quot;Monday Tuesday Wednesday Thursday Friday Saturday Sunday&quot;, 1)
    
        workdays = WD.irange(WD.Monday,WD.Friday)    # inclusive ranges are better
                                                     # for Enums, hence &quot;irange&quot;
                                                     # equivalent: workdays = WD[1:6]
        print &quot;workdays: &quot;, workdays
        print &quot;WD[1:6]: &quot;, WD[1:6]
    

    (comment continued...)

  5. 5. At 1:58 p.m. on 21 mar 2005, Martin Miller said:

    (...continued from previous comment)

        for i in WD.each():
            if i in workdays:
                print i.asInt, i.asString + &quot; is a work day&quot;
            else:
                print i.asInt, i.asString + &quot; is a weekend day.&quot;
    
        print &quot;There are&quot;, len(WD), &quot;days per week&quot;
        if &quot;Monday&quot; in WD: print &quot;Monday is a valid name&quot;
        if not &quot;August&quot; in WD: print &quot;August is not&quot;
    
        # some systematic examples:
    
        print &quot;WD.Monday.asString:&quot;, WD.Monday.asString
        print &quot;WD.Monday.asInt:&quot;, WD.Monday.asInt
        print &quot;WD.stringToInt('Monday'):&quot;, WD.stringToInt('Monday')
        print &quot;WD[2].asString:&quot;, WD[2].asString
        print &quot;WD.intToString(2):&quot;, WD.intToString(2)
        print &quot;WD['Tuesday'].asInt:&quot;, WD['Tuesday'].asInt
        print &quot;WD.Saturday &gt; WD.Tuesday:&quot;, WD.Saturday &gt; WD.Tuesday
        print &quot;'Monday' in WD:&quot;, &quot;Monday&quot; in WD
        print &quot;6 in WD:&quot;, 6 in WD
        print &quot;0 in WD:&quot;, 0 in WD
        print &quot;WD.each():&quot;, WD.each()
        print &quot;WD.eachString():&quot;, WD.eachString()
    
  6. 6. At 7:56 a.m. on 21 feb 2007, Will Ware said:

    Specified values should be considered before unspecified. 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]
    

Sign in to comment