Search code examples
pythonconventionscontrol-flow

Most Pythonic way to select behavior based on Type of input?


I have some functions that have implementation details that depend on which type of object is passed to them (specifically, it's to pick the proper method to link Django models to generate QuerySets). Which of the two following options is the more Pythonic way to implement things?

If ladders

def do_something(thing: SuperClass) -> "QuerySet[SomethingElse]":
    if isinstance(thing, SubClassA):
        return thing.property_set.all()
    if isinstance(thing, SubClassB):
        return thing.method()
    if isinstance(thing, SubClassC):
        return a_function(thing)
    if isinstance(thing, SubClassD):
        return SomethingElse.objects.filter(some_property__in=thing.another_property_set.all())
    return SomethingElse.objects.none()

Dictionary

def do_something(thing: SuperClass) -> "QuerySet[SomethingElse]":
    return {
        SubClassA: thing.property_set.all(),
        SubClassB: thing.method(),
        SubClassC: a_function(thing),
        SubClassD: SomethingElse.objects.filter(some_property__in=thing.another_property_set.all()),
    }.get(type(thing), SomethingElse.objects.none())

The dictionary option has less repeated code and fewer lines but the if ladders make PyCharm & MyPy happier (especially with type-checking).

I assume that any performance difference between the two would be negligible unless it's in an inner loop of a frequently-called routine (as in >>1 request/second).


Solution

  • This is exactly the type of problem polymorphism aims to solve, and the "Pythonic" way to solve this problem is to use polymorphism. Following the notion to "encapsulate what varies", I'd recommend creating a base "interface" that all classes implement, then just call a method of the same name on all classes.

    I put "interface" in quotation marks, because Python doesn't really have interfaces as they're commonly known in OOP. So, you'll have to make do with using subclasses, and enforcing the method signature manually (i.e. by being careful).

    To demonstrate:

    class SuperClass:
    
        # define the method signature here (mostly for documentation purposes)
        def do_something(self):
            pass
    
    class SubClassA(SuperClass):
    
        # Be careful to override this method with the same signature as shown in
        # SuperClass. (In this case, there aren't any arguments.)
        def do_something(self):
            print("Override A")
    
    class SubClassB(SuperClass):
    
        def do_something(self):
            print("Override B")
    
    if __name__ == '__main__':
        import random
    
        a = SubClassA()
        b = SubClassB()
    
        chosen = random.choice([a, b])
    
        # We don't have to worry about which subclass was chosen, because they
        # share the same interface. That is, we _know_ there will be a
        # `do_something` method on it that takes no arguments.
        chosen.do_something()