Enums in Python

Posted on Fri, 17 Feb 2017 in Python

Enums are less used feature in Python. We as programmers prefer to use weird dicts or lists where enum should be used. Why? Because it is rather new feature and you have to install back port library if you use Python 2.7. However, it's better to use them in many cases.

It is easy to find in any Python code several checks against string literals. We use such things to detect post types, client roles and so on. This code usually tends to change in time. We change set of used literals. That makes such code hard to test and maintain.

For example, we have such code.

if a.type == 'article':
    do_something_if_article(a)
elif a.type == 'review':
    do_something_if_review(a)
else:
    do_something()

Is it OK? Not really. What if we decided to change type name from "review" to "reviews". In this case, instead of executing do_something_if_review function we'll execute do_something. Even worse this code isn't easy to refactor. Name of the type could be used in many situations and variants. Could we make this code better? Sure!

ARTICLE_TYPE = 'article'
REVIEW_TYPE = 'review'

if a.type == ARTICLE_TYPE:
    do_something_if_article(a)
elif a.type == REVIEW_TYPE:
    do_something_if_review(a)
else:
    do_something()

Now we use constants. It's much better. Know we have only one place where we have to change. What if we have to add new type? We should add another elif-branch. Not so cool, is it?

This code could be improved using enums:

class PostType(Enum):
    ARTICLE = do_something_if_article
    REVIEW = do_something_if_review
    OTHER = do_something

    @classmethod
    def of(cls, type_name):
        return getattr(cls, type_name.upper(), cls.OTHER)

    def do_stuff(self):
        return self.value()

PostType.of(a.type).do_stuff()

It looks nice and this code is very flexible. I'm sure you have many places in your projects that you can improve using enums. Let's do it!

P.S. pip install enum34 for Python 2.7