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.
![[info]](https://l-stat.livejournal.com/img/userinfo.gif){kind=small}
