Choosing subclass by parameter in Python

Let’s delve into a pretty common situation. You have a number of subclasses and you have to pick one according to some parameter provided by a user (usually called type or something like this).

Here is an example: your user sends a message via your service specifying type of the message and some extra parameters.

This code works quite well but there is a number of techniques you can use to improve it.

Factory

First of all, I want a separate entity to care about message objects creation. There are two conventional types of that entity: factory or factory method.

Factory example:

Factory method example:

For me, factory is kinda redundant here, so I stick to the factory method approach in the following examples.

Choose class itself

Now we improve the code of factory method itself. For instance, since classes are first class objects in Python, you don’t really need to choose a message object, just choose class:

Use mapping

Now it’s obvious that we deal with a plain mapping types to classes. So it may be a good idea to transform our if statements to an actual map:

Note, that we can’t make MESSAGE_TYPE_TO_CLASS_MAP a class attribute since descendants of BaseMessage should be declared after BaseMessage and therefor not yet exist in the time of BaseMessage creation.

Generate class name (don’t)

You may notice that the name of every class almost matches the name of the type. We can use that to generate class name by type name and don’t store an actual mapping:

But you actually shouldn't do this. There are a number of problems with this approach. First of all, you can’t guarantee that some class called SomethingMessage will be found despite of the fact you actually don’t want it to. Second, it’s not as obvious as it can be. Third, your IDE probably won’t understand a thing in this mess.

If for some reason you decide to stick to this method, you should at least check whether the found class is a subclass of BaseMessage.

Register subclasses

The plain mapping we use is still not perfect. If you create a new message subclass you have to both define it and add to the mapping. It would be a good idea to mark the subclass somehow so it will appear in the mapping automatically. And this is where python decorators shine.

Remember I said we can’t make the mapping a BaseMessage attribute because descendant are not yet defined? Now we can make them to add themselves to the mapping upon declaration.

For this we need subclasses class attribute and register_subclass decorator to add a subclass to that attribute. Note, that register_subclass is a parameterized decorator so it returns the function that accepts a subclass as an argument.

This approach have been working for me pretty well for a long time. The best thing about it is every time new programmer create a new message subclass, he just copies the existing one as a scaffold and never forgets to add the decorator.