Back to all posts

Avoiding long if-else statements in Python

When writing Python code you might end up with something like this:

def get_device_os(device_name: str):
  if device_name == "iPhone 15 Pro":
    return "iOS"
  elif device_name == "iPhone 12 Pro Max":
    return "iOS"
  elif device_name == "Google Pixel 8 Pro":
    return "Android"
  elif device_name == "Samsung Galaxy S23":
    return "Android"
  else:
    return "Unknown"

So why is this problematic?

  • When you want to add a new device name you need to stretch the if-else statement
  • The if-else statement is static. You cannot fetch the names from a dynamic source such as a database
  • You cannot re-use the translation from device name to operating system outside of the function

How can we make this more generic and re-usable? Since the given if-else statement is basically just a mapping from device name to operating system Python's dict data structure might come in handy here:

# Can be stored in e.g. a central config file
DEVICE_OS_MAPPING: dict[str, str] = {
  "iPhone 15 Pro": "iOS",
  "iPhone 12 Pro Max", "iOS",
  "Google Pixel 8 Pro": "Android",
  "Samsung Galaxy S23": "Android"
}

def get_device_os(device_name: str) -> str:
  return DEVICE_OS_MAPPING.get(device_name, "Unknown")

In this example, we replaced the static if-else statement by a re-usable configurable dictionary.

Bonus #1

We can also use the dictionary for more than just translating a device name to an operating system name.

def get_all_known_devices() -> list[str]:
  return list(DEVICE_OS_MAPPING.keys())

def get_all_known_os() -> list[str]:
  return list(DEVICE_OS_MAPPING.values())

Bonus #2

Since the operating system is a well-defined collection of constants we can be more explicit about this by using an Enum:

from enum import Enum

class OS(Enum):
  IOS = "iOS"
  ANDROID = "Android"
  UNKNOWN = "Unknown"

DEVICE_OS_MAPPING: dict[str, OS] = {
  "iPhone 15 Pro": OS.IOS,
  "iPhone 12 Pro Max", OS.IOS,
  "Google Pixel 8 Pro": OS.ANDROID,
  "Samsung Galaxy S23": OS.ANDROID
}

def get_device_os(device_name: str) -> OS:
  return DEVICE_OS_MAPPING.get(device_name, OS.UNKNOWN)

def get_all_known_os() -> list[OS]:
  return list(DEVICE_OS_MAPPING.values())

Note how the type hints clearly indicate what kind of string is returned here. A user of these functions can now easily identify what kind of values can be returned.