Previous Entry Share
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)
<class '__main__.Status'>
>>> 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.

  • 1
Holy fuckballs, it's a Ch'marr!

Apparently I'm alive, against all reason!

Against all lacking signs thereof, how've you been old bean?

Its been somewhat of a roller coaster... and in the process of deciding if I want to ride it again, now :)

Are you sure? We should get a second opinion on that. You can never trust zombies you know.

Why - hello there! :D Now this is a post in my List, I wasn't expecting...

I thought you died in 'Nam or something...

I think I might have!

  • 1
?

Log in

No account? Create an account