Python Enum on steroids

We all do love and remember enum type from C. And I have always been missing it for Python. You can get along with class definitions for starters:

class UserRoles(object):
    root = 10
    user = 20
    nobody = 30

But it quickly complicates:

def check_perms(role):
    if role != UserRoles.root:
        raise Exception("Access denied for group %s" % role)

The above will print something like Access denied for group 20 which is not very user friendly.

Thus over time, I've came to the following enum requirements that work for me:

• Enum value should give easy access to its name
• Enums should behave like native types they where assigned to. For example, integer enums should not differ from integers, when storing them in SQL datatabase
• I should be able to assign descriptions to enums. I.e. print UserRoles.root.desc
• I should be able to resolve enum by its name. I.e. role = UserRoles.__by_name__('admin')
• I should be able to cast native type to enum. I.e. role = UserRoles(10)

So with a little help of metaclasses, I've managed to create a very useful enum class for myself. Lets see it in action first:

class UserRoles(Enum):
    root = 10
    user = 20
    user__desc = "Regular users"
    nobody = 30

>>> UserRoles.root
10
>>> UserRoles.root.name
'root'  # Wow, I have a name!
>>> UserRoles.user.desc
'Regular users'   # And even description 
>>> UserRoles.root == 10 
True    # I look and quack just like the native type I was assigned to
>>> role = UserRoles(10)   # Cast me! (with the value you get from DB for example)
>>> role
10
>>> role.name
'root'
>>> role == 10
True
>>> role == UserRoles.root
True
>>> role = UserRoles.__by_name__('root')    # Get me by name
>>> role
10
>>> role in UserRoles    # You can iterate us
True
>>> for role in UserRoles:
...   print "Role %s evaluates to %s" % (role.name, role)
... 
Role root evaluates to 10
Role user evaluates to 20
Role nobody evaluates to 30
>>> len(UserRoles)    # We are massive! :)
3
>>> 

Now if you like the above goodness, here is the implementations of Enum

class EnumMeta(type):
    __excludes__ = ()

    def __new__(mcs, name, base, d):
        cls = super(EnumMeta, mcs).__new__(mcs, name, base, d)
        cls.__values__ = set()
        cls.__members__ = dict()
        for k, v in cls.__dict__.items():
            if k in cls.__excludes__ or k.startswith("_") or k.find("__") > 0:
                continue

            pname = getattr(cls, k+"__name", k)
            pdesc = getattr(cls, k+"__desc", "")
            prop = type(type(v).__name__, (type(v),), {"name" : pname, "desc" : pdesc})
            p = prop(v)
            setattr(cls, k, p)
            cls.__values__.add(p)
            cls.__members__[v] = p

        return cls

    def __contains__(self, val):
        return val in self.__values__

    def __iter__(self):
        for i in self.__values__:
            yield i

    def __len__(self):
        return len(self.__values__)

    def __by_name__(self, name):
        if not hasattr(self, name):
            raise ValueError("%s has no property named %s" % (self.__name__, name))
        return getattr(self, name)

    def __call__(self, val):
        if val not in self.__members__:
            raise ValueError("%s has no property with value %s" % (self.__name__, val))
        return self.__members__[val]
                   
                   
class Enum(object):       
    __metaclass__ = EnumMeta  

One can notice that all of the heavy work is done during creation of enum class (i.e. when you import your code). After that you working with native types derivatives, thus performance penalty compared to using purely unorganized native types is mostly negligible. Here are some more explanations to the code:

• Enum has __metaclass__ set. Metaclass is "a class of a class". In classic OOP we have Class → Object. In Python there is one more level: Type → Class → Object. Object is instance of its Class, similarly Class is instance of its Type
• The __new__ method - when it belongs to a class, its called when a new instance is created. But when it belongs to metaclass, its called when class is created - Enum (sub)class in our case
• So once module containing a class that inherits fromEnum is imported, we drop into process of new class creation
• We inspect this class, iterating over all its members that do not start with _ and do not have have __
• Now the magic type-type-type line - here what it actually does:
  • Identify a class of the defined property and create a new class that inherits from it
  • Then create new instance of this class based on original property (copy-constructor sort of)
  • On the go, we assign this class name and description attributes
• To spice things up, the metaclass also defines all of the regular Python object model methods to make newly created enum class callable, iteratable, etc.