Archive for January, 2012

Creating an “enum” in Python

January 26, 2012

In a project I’ve been playing with, I have been using a set of global constants to define types of message that can be sent between clients. In an effort to clean up the code a little bit and streamline all the constants (and avoid collisions, since they are defined in multiple modules) I moved them to their own module (like in this interesting post).

That’s a fine solution, but I was looking for a few more features. The biggest feature I wanted was to be able to throw them in “”.format() and have a human readable name instead of a number, which is a big help when debugging. It would also be nice if I could use isinstance to check types of constants. Essentially, I wanted to mimic the behavior found in pygtk’s constants:

[~]|1> from gi.repository import Gtk
[~]|3> type(Gtk.MessageType)
   <3> <type 'type'>
[~]|4> type(Gtk.MessageType.WARNING)
   <4> <class 'gi.repository.Gtk.GtkMessageType'>
[~]|5> repr(Gtk.MessageType.WARNING)
   <5> '<enum GTK_MESSAGE_WARNING of type GtkMessageType>'
[~]|6> isinstance(Gtk.MessageType.WARNING, Gtk.MessageType)
   <6> True

This behavior sounds a bit like what you get from ‘enum’s in most languages, but, since there are no enums in Python (everthing is an object), I started googling for a clever way to fake it. It’s pretty easy to find lots of clever Python tricks if you search around.

I ran across this post at StackOverflow: http://stackoverflow.com/questions/36932/whats-the-best-way-to-implement-an-enum-in-python

The current top solution is simply a class with constants:

class Animal:
    DOG=1
    CAT=2

x = Animal.DOG

This is essentially the same as the previously linked post, replacing module with class. You can even stick your class in a module if you wanted to access it from multiple modules.

The second solution was pretty clever and had a lot of votes:

def enum(**enums):
    return type('Enum', (), enums)

>>> Numbers = enum(ONE=1, TWO=2, THREE='three')
>>> Numbers.ONE
1
>>> Numbers.TWO
2
>>> Numbers.THREE
'three'

This looks like magic, but it’s just a dynamic version of the previous solution. The function uses the built-in type() function to dynamically create a class of constants. It’s neat to be able to create it dynamically at run time instead of hard coding it, but it still doesn’t have any of the behavior of Gtk.MessageType. Because it dynamically creates a class, though, we can add extra behavior to it:

def enum_type(**enums):
    '''simple enums with a type'''
    class Enum(object):
        def __new__(cls, val):
            o = object.__new__(cls)
            o.value = val
            return o
        def __call__(self):
            return self.value

    for key,val in enums.items():
        setattr(Enum, key, Enum(val))

    return Enum

def enum_base(t, **enums):
    '''enums with a base class'''
    T = type('Enum', (t,), {})
    for key,val in enums.items():
        setattr(T, key, T(val))

    return T

The first function creates a generic enum class that can store any type of value. The second, and simpler, function inherits a base type (like int). The second class will create something very similar to Gtk.MessageType:

[~/code/python]|16> T = enums.enum_base(int, one=1,two=2,three=3)
[~/code/python]|17> x = T.two
[~/code/python]|18> x
               <18> 2
[~/code/python]|19> repr(x)
               <19> '2'
[~/code/python]|20> type(x)
               <20> <class 'enums.Enum'>
[~/code/python]|21> isinstance(x,T)
               <21> True
[~/code/python]|22> x + 5
               <22> 7

Neat! The returned class behaves (almost) exactly like I wanted. I can add extra behavior by overriding magic class methods and adding them to the third argument to type() (or I could just define the class in the function). After adding a new __repr__, the ability to name the “Enum” class, and the ability to add new values to the class, I end up with:

def enum(name, _type, *lst, **enums):
    '''
        Dynamically create enum-like class

        :param name: name of the class

        :param _type: inherited base class (like int)

        :param *lst: list of names to enumerate (ie: ONE, TWO)

        :param **enums: dict enumerations (ie: ONE=1,TWO=2)
    '''
    def _new(cls, k, v):
        obj = super(T, cls).__new__(cls, v)
        obj._name = k
        return obj

    def _repr(self):
        return ''.format(self._name, name, _type.__name__, _type(self))

    @staticmethod
    def add(*lst, **enums):
        vals = list(T._enums.keys())
        for key,val in enums.items():
            if val in vals:
                raise ValueError, "{0}'s value {1} already assigned to {2}"\
                                .format(key, val, T._enums[val])
            T._enums[val] = key
            setattr(T, key, T(key,val))
            vals.append(val)
        mx = max(vals+[0,])
        for key in lst:
            val = mx+1
            T._enums[val] = key
            setattr(T, key, T(key,val))
            vals.append(val)
            mx = val

    T = type(name, (_type,), {'__new__':_new,
                              '__repr__':_repr,
                              'add':add})

    T._enums = {}
    T.add(*lst, **enums)

    return T

Ok, it’s starting to look a little more complicated, but I added some extra ‘fluff’ to make it easy to check what has already been defined and to avoid value collisions. There is one more feature I wanted to add. In my project I pack/unpack these values into data with struct.pack/struct.unpack. I would like the ability to ‘cast’ or convert the unpacked integers back into the Enum type, the way int(‘1’) == 1.

As it turns out, you can’t just add a function to the Enum class to get this behavior, because it’s handled in type()’s __call__ function. Metaclass time! A simple metaclass that only extends __call__ and inherits type will add the desired behavior:

def enum(name, _type, *lst, **enums):
    '''
        Dynamically create enum-like class

        :param name: name of the class

        :param _type: inherited base class (like int)

        :param *lst: list of names to enumerate (ie: ONE, TWO)

        :param **enums: dict enumerations (ie: ONE=1,TWO=2)
    '''

    class Type(type):
        '''
            metaclass for new enum type, to support casting
        '''
        def __call__(cls, *args):
            if len(args) > 1:
                return super(Type, cls).__call__(*args)
            else:
                x = args[0]
                if isinstance(x, str):
                    if x in T._enums.values():
                        return getattr(T, x)
                    else:
                        return _type(x)
                elif isinstance(x, _type):
                    return getattr(T, T._enums[x])
                else:
                    raise TypeError("invalid argument type, must be str or {0}"
                                        .format(_type.__name__))

    def _new(cls, k, v):
        obj = super(T, cls).__new__(cls, v)
        obj._name = k
        return obj

    def _str(self):
        return self._name

    def _repr(self):
        return ''.format(self._name, name,
                                                _type.__name__, _type(self))

    @staticmethod
    def add(*lst, **enums):
        vals = list(T._enums.keys())
        for key,val in enums.items():
            if val in vals:
                raise ValueError, "{0}'s value {1} already assigned to {2}"\
                                .format(key, val, T._enums[val])
            T._enums[val] = key
            setattr(T, key, T(key,val))
            vals.append(val)
        mx = max(vals+[0,])
        for key in lst:
            val = mx+1
            T._enums[val] = key
            setattr(T, key, T(key,val))
            vals.append(val)
            mx = val

    T = Type(name, (_type,), {'__new__':_new,
#                              '__metaclass__':Meta,
                              '__str__':_str,
                              '__repr__':_repr,
                              'add':add})

    T._enums = {}
    T.add(*lst, **enums)

    return T
[~/code/python]|24> T = enums.enum('MyType',int,one=1,two=2,three=3)
[~/code/python]|25> T(2)
               <25> <enum two=2 of type MyType(int)>
[~/code/python]|26> T('two')
               <26> <enum two=2 of type MyType(int)>

Beautiful!

Advertisements