Skip to main content

How to make code more readable with value dispatcher pattern

How often have you written code that looks something like this?

def big_if_function(arg):
 if arg == "A":
     # large block of code
 elif arg == "B":
     # large block of code
     if arg.startswith('A'):
         return
    # continue with large block of code
 elif arg == "C":
     # large block of code
 elif arg.endswith('Z'):
    # large block of code
 else:
     # some default large block of code
     return "default"

The consequence of this code is that it is hard to read and understand. As the number of cases increases, it becomes hard to read. If the cases are anything other than simple comparisons, the whole thing quickly becomes difficult to reason about.

Solution

Dictionary of function:

def do_something_A(arg):  
 pass  
def do_something_B(arg):  
 if arg.startswith('A'): 
     return
     # continue with large block of code pass  
def do_something_C(arg):  
 pass  
def do_something_default(arg):  
 pass  
switcher = {  
 'A': do_something_A, 
 'B': do_something_B, 
 'C': do_something_C
}  
arg = 'A'  
switcher.get(arg, do_something_default)(arg)  
  • Pros:
    • Easy to read and understand
  • Cons:
    • Need to declare dictionary in function scope everytime
    • Cannot solve the case where there are multiple matches like arg ends with ‘A’

Pattern matching:

def do_something(arg):  
 print("do something default")  

def do_something_A(arg):  
    print("do_something_A")  

def do_something_B(arg):  
    print("do_something_B")

def do_something_endswithZ(arg):  
    print("do_something_endswithZ")  

def do_something_C(arg):  
    print("do_something_C")  

def main_logic(arg):  
    match arg: 
        case "A": 
            do_something_A(arg) 
        case "B": 
            do_something_B(arg) 
        case "C": 
            do_something_C(arg) 
        case x if x.endswith("Z"): 
            do_something_endswithZ(arg) 
        case _: 
            do_something(arg)  
  • Pros:
    • Can solve the case where there are multiple matches like arg ends with ‘A’
  • Cons:
    • Need python 3.10 or above
    • Still lack the ability to find out related cases.

Value dispatcher

def value_dispatch(func):
    """
    Transforms a function into a function, that dispatches its calls based on the
    value of the first argument
    """
    funcname = getattr(func, '__name__')
    static_registry = {}
    func_registry = []

    def dispatch(arg):
        """
            Return the function that matches the argument

            Match logic:
            - Check argument with static value
            - If it does not existed, check with callable list
            - Otherwise return the default function
        """
        static_func = static_registry.get(arg)
        if not static_func:
            for cond_func, f in func_registry:
                if cond_func(arg):
                    return f
            return func
        else:
            return static_func

    def register(arg):
        """
        :param arg: can be list, callable, int, etc...

        Note:
          - Callable return must be bool
          - Callable register ordering is matter. Register carefully
        """

        def wrapper(func):
            """register a function"""
            if callable(arg):
                func_registry.append((arg, func))
            elif isinstance(arg, list):
                for a in arg:
                    static_registry[a] = func
            else:
                static_registry[arg] = func
            return func
        return wrapper

    def wrapper(*args, **kwargs):
        if not args:
            raise ValueError(f'{funcname} requires at least 1 positional argument')
        return dispatch(args[0])(*args, **kwargs)

    wrapper.register = register
    wrapper.dispatch = dispatch
    wrapper.registry = static_registry
    return wrapper

@value_dispatch
def do_something(arg):
  print("do something default")

@do_something.register('A')
def do_something_A(arg):
  print("do_something_A")

@do_something.register('B')
def do_something_B(arg):
  print("do_something_B")

@do_something.register('C')
def do_something_C(arg):
  print("do_something_C")

@do_something.register(['D', 'E', 'F'])
def do_something_D_E_F(arg):
  print("do_something_D_E_F")

@do_something.register(lambda x: x.endswith('Z'))
def do_something_endswithZ(arg):
  print("do_something ends with Z")


do_something('A')
do_something('B')
do_something('C')
do_something('D')
do_something('E')
do_something('F')
do_something('AZ')
do_something('H')

Note:

  • Currently, dispatch function is not able to handle class method :(

Comments