Python code snippet - creating enums in python

[chmarr]

There's a ton of "enum" implementations for python, using named tuples, or a special class, or other things. I wanted something that looked as much as possible as an integer, but it was easy to create new instances using a symbolic name, have it print out (repr format) with that name, but otherwise be treated like any other integer.

Here's my implementation, which uses the "type" function to create a new type, inherited from "int", with all the functionality one would expect of a decently-functional enumerated type

import inspect
def enum ( name, print_str_as_name=False, **kwargs ):
def strfn ( self ):
try:
return self._rev_mapping[self.real]
except KeyError:
return str(int(self))
def reprfn ( self ):
try:
return "%s.%s" % ( name, self._rev_mapping[self.real] )
except KeyError:
return "%s(%d)" % ( name, int(self) )
revmapping = dict ( [ (v,k) for (k,v) in kwargs.iteritems() ] )
newenum = type ( name, (int,), dict ( _mapping=kwargs, _rev_mapping=revmapping, _as_name=strfn, __repr__=reprfn ) )
# We want instances of this type to appear to be created in the caller of this function, so
# we get the calling frame's module name, and stick it on to our new type
try:
frame = inspect.currentframe()
newenum.__module__ = frame.f_back.f_globals['__name__']
finally:
del frame
if print_str_as_name:
newenum.__str__ = strfn
for k, v in kwargs.iteritems ():
setattr ( newenum, k, newenum(v) )
return newenum

And Here's some example usage:

>>> Status = enum ( "Status", waiting=0, running=1, complete=2 )
>>> s = Status.running
>>> type(s)

>>> s
Status.running
>>> print "Status is", s
Status is 1
>>> print "Status is", s._as_name()
Status is running
>>> print "Status: %r" % s
Status: Status.running
>>> s == Status.running
True
>>> s == 1
True
>>> s == Status.complete
False
>>> int(s)
1

Turning an integer back into the enum type is easy.

>>> s = Status(2)
>>> s
Status.complete

If you would rather the "str" representation to show the name, rather than the integer, create the enum as follows. Note that this may break applications that expects to have the "str" turn out as an integer. Eg: web frameworks.

>>> Status = enum ( "Status", print_str_as_name=True, waiting=0, running=1, complete=2 )
>>> s = Status.running
>>> print "Status is", s
Status is running

And it pickles efficiently, since all instances of the enum are really references to the pre-constructed enum instances, which themselves are integers. There's a one-time overhead for the enum (about 60 bytes), then a smaller one-time overhead for each unique value of the enum that's pickled (about 20 bytes), then just 4 bytes for each object (same as a bare int).

>>> s1 = Status.waiting
>>> s2 = Status.running
>>> s3 = Status.complete
>>> pickle.dumps ( (s1,) )
'(ccopy_reg\n_reconstructor\np0\n(c__main__\nStatus\np1\nc__builtin__\nint\np2\nI0\ntp3\nRp4\ntp5\n.'
>>> pickle.dumps ( (s1,s2,s3,s1,s2,s3,s1,s2,s3) )
'(ccopy_reg\n_reconstructor\np0\n(c__main__\nStatus\np1\nc__builtin__\nint\np2\nI0\ntp3\nRp4\ng0\n(g1\ng2\nI1\ntp5\nRp6\ng0\n(g1\ng2\nI2\ntp7\nRp8\ng4\ng6\ng8\ng4\ng6\ng8\ntp9\n.'

Edit: added frame introspection, so object pickling will work correctly. Note that any enum type must be declared outside of any class definition, which is a restriction on the pickling process.