How to simplify long if statements in Python

If statements are great tools, but as the number of branches grows, they quickly become unwieldy. In this article, we will explore two alternatives to if statements:

• "match...case", a new Python 3.10 language feature.
• Dictionaries, particularly useful when dealing with user input.

This article's scenario: a long if statement

Imagine you have an if statement like this one (assume the average function has already been defined):

numbers = [1, 4, 16, 20]
action = input(f"What would you like to do with {numbers}?")  # e.g. add

if action == "add":
    print(sum(numbers))
elif action == "avg":
    print(average(numbers))
elif action == "max":
    print(max(numbers))
else:
    print("Action not recognized")

Using match...case to simplify an if statement chain

In "match...case", we still need to tell Python what the different options are, and what to do in each case:

def average(seq):
	return sum(seq) / len(seq)

numbers = [1, 4, 16, 20]
action = input(f"What would you like to do with {numbers}? ")

match action:
	case "add":
		print(sum(numbers))
	case "avg":
		print(average(numbers))
	case "max":
		print(max(numbers))
	case _:
		print("Operation not recognized.")

While this looks a bit better, there is still much of the same duplication. Each conditional branch has the print() function duplicated, and there are a lot of keywords.

Overall, the length of the code block is the same.

Plus the biggest problem of all remains: that as you add more options, the branching conditional will grow too.

Use a dictionary to simplify a long if statement

Instead of the log if-elif chain or a long match-case chain, you could store the user's options in a dictionary:

options = {
    "add": sum,
    "avg": average,
    "max": max
}

Then you could ask the user for which of the options (from the dictionary) they'd like to use:

options = {
    "add": sum,
    "avg": average,
    "max": max
}
numbers = [1, 4, 16, 20]

action = input(f"What would you like to do with {numbers}?")  # e.g. add

With this, we can retrieve the function from the dictionary directly:

options = {
    "add": sum,
    "avg": average,
    "max": max
}
numbers = [1, 4, 16, 20]

action = input(f"What would you like to do with {numbers}?")  # e.g. add

operation = options.get(action)

Since the dictionary maps strings to functions, the operation variable would now contain the function we want to run.

All that's left is to run operation(numbers) to get our result. If the user entered 'add', then operation will be the sum function.

We should also do some error checking, to make sure we don't try to run a function that doesn't exist if the user entered something that isn't one of the dictionary's keys.

options = {
    "add": sum,
    "avg": average,
    "max": max
}
numbers = [1, 4, 16, 20]

action = input(f"What would you like to do with {numbers}?")  # e.g. add

operation = options.get(action)

if operation:
    operation(numbers)
else:
    print("Action not recognized")

You still need the one if statement just in case the user chooses something that doesn't have a key in the dictionary, but this is a single-branch if statement that won't grow over time.

Another benefit is that you can easily tell the user which options are available to them by using the dictionary keys:

option_texts = '|'.join(options.keys()
action = input(f"What would you like to do with {numbers}? ({option_texts}) ")
# Would show "What would you like to do with [1, 4, 16, 20]? (add|avg|max) "

Conclusion

In this post we've seen two ways you can simplify if statements: by using "match...case" and by using dictionaries.

If you want to learn more about Python, consider enrolling in our Complete Python Course which takes you from beginner all the way to advanced (including OOP, web development, async development, and much more!). We have a 30-day money-back guarantee, so you really have nothing to lose by giving it a try. We'd love to have you!

Photo by Christin Hume on Unsplash