Fighting with type hints – overloaded decorator

How to properly type-annotate the following decorator?

def pass_thru_none(func):
    def wrapper(value):
        if value is None:
            return None
        return func(value)

    return wrapper

The solution is not simple at all. Looking for it, we will explore many advanced Python features:

• decorators,
• Callable type,
• TypeVar and Generic,
• the @overload decorator,
• and the __call__ protocol.

Example problem

Consider that we have a Python applications that reads data from a database and returns them to a JSON API.

We have classes for both the database models and API schemas. In practice, the classes could be SQLAlchemy and pydantic models, respectively, but for the sake of simplicity, I will use dataclasses in examples.

import datetime as dt
from dataclasses import dataclass
from typing import Optional


def format_datetime(value: dt.datetime) -> str:
    return value.isoformat()


@dataclass
class UserModel:

    name: str
    create_date: dt.datetime
    login_date: Optional[dt.datetime]


@dataclass
class UserSchema:

    name: str
    createDate: str
    loginDate: Optional[str]

    @staticmethod
    def from_model(model: UserModel) -> "UserSchema":
        return UserSchema(
            name=model.name,
            createDate=format_datetime(model.create_date),
            loginDate=format_datetime(model.login_date)
            if model.login_date is not None
            else None,
        )

In the example, UserModel could be a database model and UserSchema would be a schema for a JSON API. UserSchema.from_model() can construct the latter from the former, renaming field names to camelCase and converting dates to strings.

Handling None values

The example has one detail that does not look pretty – the ternary operator needed to take care of login_date, which is optional. Real-world example would have dozens of checks like this, so we may decide to improve the format_datetime() function to take care of that:

import datetime as dt
from typing import Optional


def format_datetime(value: Optional[dt.datetime]) -> Optional[str]:
    if value is None:
        return None
    return value.isoformat()

Unfortunately, Mypy complains about the improved function:

error: Argument "createDate" to "UserSchema" has incompatible type "Optional[str]"; expected "str" [arg-type]

The problem is that Mypy cannot infer that format_datetime(model.create_date) never returns None.

We have to use the @overload decorator to provide the typechecker with more details:

import datetime as dt
from typing import Optional, overload

@overload
def format_datetime(value: None) -> None:
    ...


@overload
def format_datetime(value: dt.datetime) -> str:
    ...


def format_datetime(value: Optional[dt.datetime]) -> Optional[str]:
    if value is None:
        return None
    return value.isoformat()

Reusable decorator

The format_datetime() function is an example of one conversion function, but a real application needs many functions like that.

It would be nice if all conversion functions passed None thru, without having to copy the logic (with two @overload variants) for each type. This is the motivation for the decorator from the beginning of this article:

import datetime as dt

def pass_thru_none(func):
    def wrapper(value):
        if value is None:
            return None
        return func(value)

    return wrapper


@pass_thru_none
def format_datetime(value: dt.datetime) -> str:
    return value.isoformat()

The decorator is straightforward. A few years ago, we would be done. But these days, the situation different. We want our type-safe code, and functions decorated by an untyped decorator would be untyped too.

To type the decorator, we will need at least two features from the typing module:

• Callable as a type of input and output of the decorator.
• TypeVar to make the decorator generic, handling conversion function for any type.

from typing import Callable, Optional, TypeVar

INPUT_T = TypeVar("INPUT_T")
OUTPUT_T = TypeVar("OUTPUT_T")


def pass_thru_none(
    func: Callable[[INPUT_T], OUTPUT_T]
) -> Callable[[Optional[INPUT_T]], Optional[OUTPUT_T]]:
    def wrapper(value: Optional[INPUT_T]) -> Optional[OUTPUT_T]:
        if value is None:
            return None
        return func(value)

    return wrapper

The code is getting ugly, yet it is not enough. We are getting an error that we saw earlier:

error: Argument "createDate" to "UserSchema" has incompatible type "Optional[str]"; expected "str" [arg-type]

Overloading the decorator

We need @overload for the decorator. But there is no syntax how to combine it with Callable.

Another complication seems unavoidable. Let's convert the decorator into a class with a __call__ method. The class has to be declared as Generic to preserve type information between __init__ and __call__.

from typing import Callable, Generic, Optional, TypeVar, overload

INPUT_T = TypeVar("INPUT_T")
OUTPUT_T = TypeVar("OUTPUT_T")


class pass_thru_none(Generic[INPUT_T, OUTPUT_T]):
    def __init__(self, func: Callable[[INPUT_T], OUTPUT_T]) -> None:
        self.func = func

    @overload
    def __call__(self, value: None) -> None:
        ...

    @overload
    def __call__(self, value: INPUT_T) -> OUTPUT_T:
        ...

    def __call__(self, value: Optional[INPUT_T]) -> Optional[OUTPUT_T]:
        if value is None:
            return None
        return self.func(value)

I know, this is getting crazy. All we wanted is to add type hints to a 6-line decorator, and we have this cryptic class. Yet we are not done yet. Mypy complains:

error: Overloaded function signatures 1 and 2 overlap with incompatible return types [misc]

Frankly, I expected all the problems so far, but this last resistance surprised me. Mypy docs explain the problem quite well – an unbound type variable can include None, so the two overloads may not be distinguishable.

A recommended fix is to swap order of the overloads. The swap alone does not help in our case, but when we relax one overload, the following example finally works:

from typing import Callable, Generic, Optional, TypeVar, overload

INPUT_T = TypeVar("INPUT_T")
OUTPUT_T = TypeVar("OUTPUT_T")


class pass_thru_none(Generic[INPUT_T, OUTPUT_T]):
    def __init__(self, func: Callable[[INPUT_T], OUTPUT_T]) -> None:
        self.func = func

    @overload
    def __call__(self, value: INPUT_T) -> OUTPUT_T:
        ...

    @overload
    def __call__(self, value: None) -> Optional[OUTPUT_T]:
        ...

    def __call__(self, value: Optional[INPUT_T]) -> Optional[OUTPUT_T]:
        if value is None:
            return None
        return self.func(value)

Because of the second overload was relaxed, Mypy can no longer infer that decorated functions always pass thru None, but that's not needed in practice.

Is it all?

The decorator with proper type annotations has grown significantly, but we might not be done yet.

We may want to update_wrapper. But it is OK to have an pass_thru_none instance with attributes of some function?

We should also consider how to use the decorator with class methods. The wrapper could accept variable *args, including self or cls at the first position. In such case, ParamSpec might be needed.

But we got too far already, so I will keep further complications for the reader.

Conclusion

Full code can be found in Gist.

Do I recommend the monster? No. Maybe, the decorator wasn't such a good idea. On the other hand, I can imagine similar code hidden deep in a framework, where proper type hinting is more important that having the most beautiful code.

The example shows how type hints in Python force us to a different code style. Some people could argue that this is not “Pythonic”. But it is a nice showcase of advanced tools from the typing module, so it was worth being nerd sniped by it :)