Tweag

Safe composable Python

6 June 2024 — by Guillaume Desforges

Writing modular code is a challenge any large project faces, and that stands even more true for Python monoliths. Python’s versatility is convenient, yet a very big gun to shoot one’s foot with. If there’s one thing I’ve learnt over the years of building software, it’s the importance of testing to mitigate risks and allow refactoring in the long run.

In my experience however, I’ve seen much less testing in projects than I would’ve liked. If it can be rarely attributed to plain malevolence or stupidity, the most recurrent reason is: it is hard and it takes too much time. The issue is, it’s often true!

So how can we lower the barrier to make testing easier?

Function composition to the rescue

Behind this seemingly scientific section title rests a very simple principle.

Let’s look at the following example:

import sdk

from model import User


def find_admin() -> User | None:
    users: list[User] = sdk.get_users()
    return next(u for u in users if u.name == "admin", None)

This function find_admin takes no arguments and returns either an object of type User for the first admin user it finds or None.

Now, what would it take to test this function find_admin?

In this example, I pretend that there is a module sdk that gives me some SDK to work with APIs. When called, sdk.get_users will make an HTTP request, parse it, and return a list[User]. Such a function that “interacts with the outside world” is often qualified as effectful.

There are a few strategies to test a function that uses an effectful function:

The first option is to set up the testing environment to correspond to what the effectful function expects. In this case, let the tests run against the real API. Hopefully you can do that against a sandbox and not your production API, but even then you will soon run into many challenges. For instance, your tests will require the sandbox to be in a specific state every time you run them, which is not only hard to implement, but might also mean that you either cannot run multiple tests in parallel, or that you need to support multiple isolated test environments. This makes testing impracticable, soon to be given up. This option is not good at all.

The second and more reasonable option is to mock the API and make your effectful code interact with this mock. In this example, at the start of a testing session, we would spawn an HTTP server or monkeypatch the SDK to emulate locally on your computer the API that the SDK tries to interact with. However, mocking an entire system with different interacting modules can quickly become cumbersome and make the code too heavyweight. We know how it ends: it’s too complicated, takes too much time, hence you give up testing. This option is not good enough either.

The third option is to completely change how your code works:

As it is, our function find_admin uses sdk.get_users. It depends on this function, and it’s the root of our complications. Let’s just get rid of it then!

def find_admin(get_users: Callable[[], list[User]]) -> User | None:
    users: list[User] = get_users()
    return next(u for u in users if u.name == "admin", None)

What happened? Instead of calling sdk.get_users directly in the body of find_admin, we put an argument get_users that the body of the function calls, with a type hint Callable[[], list[User]] that indicates explicitly that it should be a function that takes no argument ([]) and returns a list[User].

As such, it is now the caller’s responsibility to inject the right function when calling find_admin.

def some_function():
    # get actual admin from the API
    admin = find_admin(get_users=sdk.get_users)

Sounds reasonable. Now, can we test find_admin?

def test_find_admin():
    def get_users():
        return [User(name="william")]

    assert find_admin(get_users=get_users) is None

No testing with the production API, no mocks, everything runs with basic Python. And all we need is a simple function that returns a list[User].1

Function composition in Python makes it easier to write modular, reusable and testable code. This style compels us to use type hints, since without it would be hard to know which kind of function should be given, there would be no type checking and no auto-completion. The type hint Callable is needed to make this way of coding usable, unfortunately it is often not enough.

Complex function typing with Protocol

Let’s consider that there is a more efficient function for our purpose, sdk.find_user(name: str, allow_admin: bool) -> User | None.

Now, let’s try to use this function for find_admin:

def find_admin(find_user: Callable[[str, bool], User | None]) -> User | None:
    return find_user(name="admin", allow_admin=True)

The type hint for find_user means that it must be a function that takes two arguments, the first a str and the second a bool, and returns either a User or None.

Huh oh, what is that?

error: Expected 2 more positional arguments

Our type checker, in my case pyright, is not pleased at all. Callable can specify the positional arguments of the function expected, but it can’t define a more precise type, for instance the names of the arguments. The type checker, at call site, can’t infer from Callable[[str, bool], User | None] that one argument has a name name. It can only check positional arguments, of which we pass none, hence the error.

We could not use keyword arguments when calling find_user, but in complex and larger applications it is more robust to pass arguments explicitly by their name. Instead, let’s use a better tool to provide a precise type hint.

This is where typing.Protocol, defined in PEP 544, is handy.

Protocol is very akin to an interface, but contrary to abc.ABC interfaces, an object does not need to inherit a protocol to type-check against it. Instead, it only needs to implement all its interfaces.2

class Flippable(Protocol):
    def flip(self: Self) -> None: ...


class Table:
    def flip(self):
        print("(╯°□°)╯︵ ┻━┻")


flippable: Flippable = Table()
flippable.flip()

You may wonder how, since there is no inheritance between Flippable and Table, the constraints of the former could apply to the latter. The key point to realize here is this line:

flippable: Flippable = Table()

This tells the type checker that the variable flippable can only be assigned objects which are compatible with the Flippable protocol, meaning objects that implement the flip method. This is the line where the type checker will compare the Table class to the Flippable protocol and decide whether the interface is properly implemented.

Let’s look at an example where we try to assign an object incompatible with the protocol.

class Duck:
    ...


other_flippable: Flippable = Duck()

If we try to assign a new Duck to a variable typed as Flippable like so, the type checker will report an error.

error: Expression of type "Duck" is incompatible with declared type "Flippable"
    "Duck" is incompatible with protocol "Flippable"
      "flip" is not present

It rightfully reports that Duck does not implement the flip method, making Duck incompatible with the type Flippable required for this variable.

Now, how can we use Protocol to type a function?

A function call is represented by the __call__ method of an object. As such, making a protocol with a signature for __call__ allows us to type a function precisely, including the names of the arguments and their types.

class FindUser(Protocol):
    def __call__(self, name: str, allow_admin: bool) -> User | None:
        ...


def find_admin(find_user: FindUser) -> User | None:
    return find_user(name="admin", allow_admin=True)

This makes our type checker happy.

The catch

Looking at the snippet of code we end up with

class FindUser(Protocol):
    def __call__(self, name: str, allow_admin: bool) -> User | None:
        ...


def find_admin(find_user: FindUser) -> User | None:
    return find_user(name="admin", allow_admin=True)


def some_processing():
    admin = find_admin(find_user=sdk.find_user)

We can anticipate a few shortcomings of this way of coding.

The only explicit reference between the protocol FindUser and the actual implementation sdk.find_user is when calling the function with the implementation find_admin(find_user=sdk.find_user). This often leads to an unsavory double-edit when changing the signature of the argument, having to change both the protocol and the implementation. Given that you will also have an implementation for testing, this will quickly turn into a triple-edit situation.

In my experience, this is only mildly annoying and can be compensated for by tooling.

Conclusion

Testing code that works with APIs, databases and any other kind of external system is hard. Composition is a great solution: not only does it make code modular and reusable, it also makes such functions with effectful code testable.

Historically, that came at the price of safety due to the limitations of Callable. That is not the case any more, as Protocol helps define precise interfaces for functions.

While focusing on the technical aspects, this way of coding is very biased towards Test Driven Development and Hexagonal Architecture. These are great principles to ship better and faster.


  1. This technique is known as ”dependency injection“.
  2. This is known as ”structural typing“.

About the author

Guillaume Desforges

Guillaume is a versatile engineer based in Paris, with fluency in machine learning, data engineering, web development and functional programming.

If you enjoyed this article, you might be interested in joining the Tweag team.

This article is licensed under a Creative Commons Attribution 4.0 International license.

Company

AboutOpen SourceCareersContact Us

Connect with us

© 2024 Modus Create, LLC

Privacy PolicySitemap