Write Python Functions Like This Or I’ll Reject Your Pull Request

This was the energy I was getting from my tech lead at work. And I actually agree with him at this point.

How we were taught to write Python functions

Here’s a simple function that takes in:

• a list of numbers num_list
• a number num
• adds num to every single number in num_list
• and returns a new list

def add_all(num_list, num):
    output = []
    for n in num_list:
        output.append(n + num)
    return output

x = add_all([3, 4, 5], 10)

print(x) # 13, 14, 15        

Problem — we don’t know at first glance:

• what data types this function takes in
• what data type this function returns

I mean, we can infer from reading the code that num_list is a list of numbers, num is a number, and that the function returns a list of numbers. But this isn’t instant.

And in a larger production-grade app, we have thousands of functions to deal with. Do we really want to spend this extra time figuring out and inferring data types?

Type Annotations (type hints) to the rescue

Python’s type annotations (or type hints) provide a way to explicitly state the expected data types of variables and return values in functions. This can make your code more readable and maintainable, especially in large codebases.

Here's how we can rewrite the previous function using type annotations:

from typing import List

def add_all(num_list: List[int], num: int) -> List[int]:
    output = []
    for n in num_list:
        output.append(n + num)
    return output

x = add_all([3, 4, 5], 10)
print(x)  # [13, 14, 15]
        

Now, at first glance, we know:

• num_list is a list of integers.
• num is an integer.
• The function returns a list of integers.

Benefits of Using Type Annotations

1. Readability: Makes the code more readable and understandable, especially for new developers joining the team.
2. Error Detection: Helps in catching errors early during development with tools like mypy.
3. IDE Support: Many modern IDEs provide better autocompletion and error checking when type annotations are used.
4. Documentation: Serves as inline documentation for your functions, reducing the need for extensive comments.

Using Mypy for Static Type Checking

To make the most out of type annotations, you can use a static type checker like mypy. It can analyze your code and ensure that the type annotations are followed correctly.

First, install mypy:

pip install mypy
        

Then, you can run mypy on your Python files:

mypy your_file.py
        

If there are any type mismatches, mypy will report them, allowing you to fix issues before they cause runtime errors.

Type Annotations for Complex Types

Type annotations aren’t just for simple types like int and str. You can also use them for more complex types such as lists, dictionaries, and even custom classes.

Example with Dictionaries

Suppose we have a function that processes a dictionary of user data:

from typing import Dict

def process_user_data(user_data: Dict[str, int]) -> None:
    for key, value in user_data.items():
        print(f"{key}: {value}")

user_data = {'Alice': 25, 'Bob': 30}
process_user_data(user_data)
        

Here, user_data is expected to be a dictionary where the keys are strings and the values are integers.

Conclusion

Using type annotations in Python is a best practice that can save time and reduce errors in the long run. It makes your codebase more robust, maintainable, and easier to understand. By incorporating type hints and tools like mypy, you can ensure that your code adheres to the expected types, leading to fewer bugs and a smoother development experience.

So, the next time you write a function, consider adding type annotations. It might take a few extra seconds, but it will pay off in the long run by making your code more readable and maintainable.