Dependency Injection (like pytest fixtures) with a Python decorator

If you’ve used Pytest, you’ll notice a neat little pattern. If you want a particular fixture, you just ask for it, seriously, check out the docs:

@pytest.fixture
def fruit_bowl():
    return [Fruit("apple"), Fruit("banana")]


def test_fruit_salad(fruit_bowl):
    # Act
    fruit_salad = FruitSalad(*fruit_bowl)

test_fruit_salad gets called when you run tests, and it will automatically pass the result of fruit bowl to test_fruit_salad. This pretty much blew my mind when I first saw it.

The following decorator helps you achieve something similar with your own functions. It lets you do something like this.

@opt_in_args
def request_a(foo: str, *, a: int):
    print(f"Foo is {foo}, a is {a}")

@opt_in_args
def request_b(*, b: str):
    print(f"b is: {b}")

@opt_in_args
def request_both(*, b: str, a: int):
    print(f"b is: {b}, a is: {a}")

request_a("Blarg")
request_b()
request_both()

Note that the * args just mean “everything after this is a keyword argument only”, it helps guard against accidentally passing in bad kwargs and you should always use it, but the above program will work the same without it. The output of the program looks like this:

Foo is Blarg, a is 123
b is: foobar
b is: foobar, a is: 123

First let’s look at the implementation of the decorator:

import inspect

def opt_in_args(func):
    def wrapper(*args, **kwargs):
        signature = inspect.signature(func)
        opted_in = {}
        optional_kwargs = {
            "a": 123,
            "b": "foobar",
        }
        for kwarg, value in optional_kwargs.items():
            if kwarg in signature.parameters:
                assert signature.parameters[kwarg].annotation == type(value)
                opted_in[kwarg] = value
        return func(*args, **kwargs, **opted_in)
    return wrapper

Let’s look at a few specific lines to illuminate what it’s doing:

def opt_in_args(func):
    def wrapper(*args, **kwargs):

This is standard decorator magic. The outer function has the function in-scope, the wrapper is (because we return it from the decorator) what actually gets called in place of the wrapped function, so we capture *args and *kwargs so we can pass them along.

signature = inspect.signature(func)

Note that we imported inspect at the top of the module. Inspect lets you introspect on python functions, and in this case we use it to grab the signature. signature.paremeters will now have a key for each parameter the function accepts; we can now write any code we like based on the params.

        optional_kwargs = {
            "a": 123,
            "b": "foobar",
        }
        for kwarg, value in optional_kwargs.items():
            if kwarg in signature.parameters:

This approach is totally optional. There are any number of ways you can enumerate the optional kwargs. You can make another decorator and supply functions which are only called if they’re requested (I think pytest fixtures must be doing something like this) or you can have a bunch of ifs to decide what to instantiate and put in.

                assert signature.parameters[kwarg].annotation == type(value)

This I just put in there to show one way you could enforce typing on the keyword arguments. This is of course also optional; you could write your decorator to only care about the name by removing this line.

                opted_in[kwarg] = value
        return func(*args, **kwargs, **opted_in)

We create a dictionary called opted_in and populate it with the additional kwargs we want to pass. The final line of the wrapper function calls the wrapped function with the args and kwargs the caller sent, unmodified, and finally with the opted_in args the function asked for in its signature.

For such a slick effect, this is a surprisingly easy decorator; decorator magic often is.