Have you ever wanted to use switch-case statement in Python code? If yes, then you have probably found out that there is no such expression in native Python syntax. Fortunately, you can implement it on your own while having some nerd-level fun. This post explains how to implement switch-case statement in Python.

TL;DR

How to implement switch-case statement in Python, like this one:

var = 666
with switch(var) as case:
    case(1, "foo")
    case(2, "yolo")
    case(3, "bar")
    case.default("dunno")
    result = case.result
print(result)
>>> "dunno"

Prerequisites & system requirements

  • can-do attitude
  • liberal approach to Python
  • time for some nerd-fun

Switch-Case idea in Python

Here’s the idea: Python has with statement, which can be used to control entering and exiting to the blocks of code. This is done by implementing __enter__ and __exit__ magic functions in your object.

Another magic function that is pretty handy is __call__. This one allows to handle object “calling”, which is basically this:

x = MyClass()
x() # <-- __call__ method of object x of type MyClass will be executed there

OK - we have our building blocks. As you’ve probably saw in the TL;DR above, our desired syntax is the following:

with switch(SOME_OBJECT) as case:
    case(1, "case value for 1")
    case(2, "case value for 2")
    case.default(r"¯\_(ツ)_/¯")
    result = case.result
print(result)

Additionaly, I would like to have option for case condition and case result to be lambda expression, like this:

y = object()
with switch(y) as case:
    case(lambda: isinstance(y, str), "Some characters")
    case(lambda: isinstance(y, int), "Numberzzzz")
    case(lambda: isinstance(y, float), "Why u not double?")
    case.default(r"¯\_(ツ)_/¯")
    result = case.result
print(result)

Switch-Case implementation in Python

Without further ado, here’s the implementation:

class switch(object):
    def __init__(self, object_to_switch_on):
        self.object_to_switch_on = object_to_switch_on
        self.default_result = None
        self.result_store = None
        self.cases = {}
        self.evaluated = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        _ = self.result

    def __call__(self, when, then):
        if when in self.cases:
            raise ValueError("This case already exists")
        else:
            self.cases[when] = then
        return self

    def default(self, then):
        self.default_result = then
        return self

    @property
    def result(self):
        if not self.evaluated:
            self.result_store = self.__evaluate__()
        return self.result_store

    def __evaluate__(self):
        from inspect import signature
        self.evaluated = True
        result_tmp = None
        if all(callable(x) for x in self.cases.keys()):
            for case_fn, case_result in self.cases.items():
                no_of_params = len(signature(case_fn).parameters)
                if no_of_params >= 2:
                    raise ValueError("Case function must have either 0 or 1 arguments")

                if no_of_params == 1 and case_fn(self.object_to_switch_on):
                    result_tmp = case_result
                    break
                elif no_of_params == 0 and case_fn():
                    result_tmp = case_result
                    break
            else:
                result_tmp = self.default_result
        elif all(not callable(x) for x in self.cases.keys()):
            result_tmp = self.cases[self.object_to_switch_on] if self.object_to_switch_on in self.cases else self.default_result
        else:
            raise ValueError("Inconsistent switch usage, use either simple objects or functions")
        if result_tmp is not None and callable(result_tmp):
            return result_tmp()
        else:
            return result_tmp

Code explanation

  1. Class name is lowercased in order to imitate basic language syntax while used.
  2. Construtor (__init__) takes “switchable” object.
  3. Magic functions: __enter__ & __exit__ are here to implement Python with scope behaviour. Exit calls result property in order to trigger the evaluation when exiting the scope (it’s here just to shorten the usage syntax with lambdas even more, see exmaples below).
  4. Magic function __call__ is used to add cases to our switch as we call it. Cases are stored internally in a dictionary for fast direct object access (as long as the case is a hashable object).
  5. default method adds default case and it’s optional to call it.
  6. result property getter is triggering switch-case evaluation. If it’s called for the first time, it invokes __evaluate__ method, which does all the magic of resolving switch-case statement result. Every subsequent call returns the result from cache.
  7. __evaluate__ method is where all of the meat is. This method checks if all of the cases are either callable objects (most likely functions or lambdas). If they are, it iterates through them and assumes that match is first match for which the function/lambda returned truthy value. If none of the cases is callable, I assume that the objects are hashable and I just lookup for their values in the internal cases dictionary. After resolving the case result, I also check, whether result is callable - if it is, I invoke it, if not I return it.

Summary

And that’s it. Don’t rate the Pythonic code ratio. Like I said, it’s only for liberal Python programmes. Hope you liked it!

Additional links & resources

Tags

python

Comments