Refactoring if else statements

Hey everyone! I have been always wondering if there are any good alternatives to if else statements since my first real programming job ever. I've seen  how my colleagues software developers used if else a lot and it worked just perfect! But at the same time some resources say that if else is a code smell and considered as bad practice.  Curiosity led me to research the topic and finally I can tell you what do I think about it.

When if else might be good

Short answer: when it's simple and testing options suit you.

Small amount of conditions is straightforward and clear. Even a beginner could understand that peace of code. Look at the code below.

# calculator.py

def calculate(operation: str, a: int, b: int) -> int:
    if operation == '+':
        return a + b
    elif operation == '*':
        return a * b
    
    raise NotImplementedError(f"{operation} not supported")

Doesn't look like as complicated code. The code structure follows the obvious patterns. Operation conditions and results are simple enough to identify what they are doing. Now let's add some tests.

from calculator import calculate


def test_multiplication():
    assert calculate(operation='*', a=2, b=2) == 4


def test_sum():
    assert calculate(operation='+', a=2, b=2) == 4

As you can see the tests are also simple enough. However the testing is wrapped around only one function as entrypoint - calculate. In order to test different logic you have to pass different parameters to this function and test the results.

There's nothing wrong if the code is simple(and not going to expand much) and we are ok with that testing goes through entrypoint as single function requiring different parameters for different branches.

The best part in it is that you can easily follow up on next elif part in your head because you've already learned the pattern the code follows on. For / it woul be a / b, for - it would a - b, etc.

# Pattern is simple enough to follow on
if operation == '+':
    return a + b
elif operation == '*':
    return a * b
elif operation == '/':
    return a / b
elif operation == '-':
    return a - b

If else conditions might be good when:

• the code is pretty small which allows to manipulate with the objects in your head;
• you can easily identify the patterns and follow up on conditions;
• you are ok with testing options;
• you are prototyping or experimenting with something;

When if else might be bad

Small amount of code is quite readable with if else statements, but when this code grows in amount you are going to face several issues:

• hard to read and remember;
• testing is not convinient if you want to test many branches;

Example of refactoring

Suppose we have two roles in the system - parent and student.
Each role must be processed in a different way. Parents provide information about children. Students provide only info about courses. So, we have to do different algorithms there.

Parent role tells to endpoint to register parent firstly and then register each of his children.
Student role tells only to register a student with specified courses he wants to enroll onto. Different logic.

def register(request):
    if request.data['role'] == 'parent':
        do_logic(request.data)
        do_another_logic()
        ... 
        register_parent(request.data)
        register_children(request.data)
        enroll_children_on_courses(request.data)
        ...
    elif request.data['role'] == 'student':
        ...
        register_student(request.data)
        ...
        enroll_student_on_courses(request.data)

Honestly speaking this is still quit readable and straightforward simple.
But in case of adding more of the roles and logic behind that we have to add another else if statement. Each time we want to change logic behind a certain role we must go to this function and scroll down to the appropriate if statement.

Someone finds this code messy. I have three approaches in my mind to refactor this. Let's investigate them.

Approach #1 - dictionary and functions

# services/registration.py

# Each role is processed in its own function
def process_parent(data)
    # Edit this function to change 
    # parent registration logic
    do_logic(data)
    do_another_logic()
    ... 
    register_parent(data)
    register_children(data)
    enroll_children_on_courses(data)
    ...


def process_student(data):
    # Edit this function to change 
    # student registration logic
    ...
    register_student(data)
    ...
    enroll_student_on_courses(data)

def register(request):
    role = request.data['role']
    
    # Straightforward pattern.
    # For each role there's only one function.
    # Easy to learn and use
    process = {
        'parent': process_parent,
        'student': process_student,
    }.get(role, process_student)
    
    process()

Dictionary restricts number of possible patterns for our brain to deal with.
We use role's name as string on the left side as key. There's always only one key, not anymore. On the right side as value we use functions. There's only one value on the right, not anymore. This allows our brain to learn one pattern and read it quickly without confusion.

Splitting into functions adds ability to test each part(or unit) individually. This is a good option if you are concerned about testing options. Now you are able to write unit tests.

# tests/unit/registration.py
from users.services.registration import process_parent, process_student

...


def test_process_parent():
    # Filling your data
    data = ...

    # Calling single function
    process_parent(data=data)
    # Checking results
    assert children_are_registered() is True
    assert children_are_onrolled(data) is True

    ...


def test_process_student():
    # Filling your data
    data = ...

    # Calling single function
    process_student(data=data)
    # Checking results
    assert student_is_registered() is True
    assert student_is_onrolled(data) is True

    ...

Looks like more structured and less messy without those if .. == .. statements. If you need to change a certain role logic - just go ahead and change appropriate function. No need to search one big function for a peace of related code.

Benefits

Simple pattern to learn

Dictionary based branching has constraints due to its inteface(1 key, 1 value). For our brain it's much easier to remember it as a pattern than if .. == .. or .. and. This results in faster processing the code by our brain and better understanding.

Easier to navigate and memorize

Splitting into functions triggers important technique used to memorize things. The technique called chunking. Chunking is a process of grouping a set of small items into a unit of information. The same relies to dividing the code into classes or modules. Such organization will be a good advantage in long term.

Unit testing

Now we are able to test each function as a unit. This allows to write much more convinient and smaller tests.

Approach #2 - OOP Style. Strategy pattern

For OOP lovers there's another approach with classes which is based on strategy pattern. For more details about the pattern you can read this resource https://refactoring.guru/design-patterns/strategy

class UserRegistration(object):
    
    def __init__(self, data):
        self.data = data
        
    def process(self):
        rainse NotImplemented


class ParentRegistration(UserRegistration):

    def process(self):
        do_logic(self.data)
        do_another_logic()
        ... 
        register_parent(self.data)
        register_children(self.data)
        enroll_children_on_courses(self.data)
        ...


class StudentRegistration(UserRegistration):
    
    def process(self)
        ...
        register_student(self.data)
        ...
        enroll_student_on_courses(self.data)


def register(request):
    role = request.data['role']

    user_registration_class = {
        'parent': ParentRegistration,
        'student': StudentRegistration,
    }.get(role, StudentRegistration)
    user_registration = user_registration_class(reques.data)
    
    user_registration.process()

Is it better than if else? Well. It might be more readable for OOP thinkers who love objects, classes, etc. For those who think about code within concepts of classes. But of course it's more complicated for beiginners who don't know much about classes and OOP concepts.

Approach #3 - decorators

strategies = dict()


def handler(role):
    def wrapper(func):
        global strategies
        strategies[role] = func

        def inner(*args, **kwargs):
            return func(*args, **kwargs)

        return inner

    return wrapper


def process_role(role, data):
    strategy = strategies.get(role)
    strategy(data)


@handler(role="parent")
def process_parent(data):
    print("Processing parent")
    print(data)
    ...


@handler(role="student")
def process_student(data):
    print("Processing student")
    print(data)
    ...


def register(request):
    process_role(request["role"], request["data"])

Personally this approach is much readable in case of adding new roles. I find it simpler and more pythonic. However, it does contains some magic behind it. The part with the decorator logic is a little bit hard to understand for beginners. The good news you don't need to return to the decorator inners if you just want to change or add new role strategies. Adding new functions in is extremly simple.

Be aware that using this approach in different files requires you to import each of the functions somewhere in your app before running to register all handlers.

UPDATE: Python 3.10 pattern matching

Since python 3.10 we have an alternative to the described methods.

def register(request):	
	match request.data['role']:
		case 'parent':
			do_logic(request.data)
            do_another_logic()
            ... 
            register_parent(request.data)
            register_children(request.data)
            enroll_children_on_courses(request.data)
            ...
		case 'student':
            ...
            register_student(self.data)
            ...
            enroll_student_on_courses(self.data)
		case role:  # Here we can capture any other roles and save them in the variable
			print(f"Role {role} can't be processed")        	

We may even use class instances and go deeper in using match case.

class Role:
    __match_args__ = ('data',)

    def __init__(self, data: dict) -> None:
        self.data = data

    def process(self) -> None:
        raise NotImplementedError


class ParentRole(Role):

    def process(self) -> None:
        print(f'Parent processing {self.data}')


class StudentRole(Role):

    def process(self) -> None
        print(f'Student processing {self.data}')


def construct_role(data: dict) -> Role:
	role = data.get('role', 'student')
    match role:
        case 'parent':
            return ParentRole(data=data)
        case 'student':
            return StudentRole(data=data)

def register(request):	
    role: Role = construct_role(request.data)

    match role:
        case ParentRole(data) as parent_role:
            print(f'Processing parent with data: {data}')
            parent_role.process()
        case StudentRole(data) as student_role:
            print(f'Processing student with data: {data}')
            student_role.process()      	

There are some of new ways for refactoring using pattern matching. Look at the example below:

my_bet = 5
match compare(my_bet, 4):  # output: 5 is greater than 4!
    case Greater(value):
        print(f'{my_bet} is greater than {value}!')
    case Less(value):
        print(f'{my_bet} is less than {value}!')
    case Equal(value):
        print(f'{my_bet} is equal than {value}!')

Look at this repo to inspect the code with compare.

Conclusion

You can refactor your if else statements using classes, decorators, functions and match case statements.
However the final choice depends on you. Choose any of the options and determine which is the best for you.