Annotating a decorator that changes the type of a keyword argument

[jwodder]

I'm trying to add type annotations to some code that, to keep things DRY, applies a decorator to a number of functions; this decorator "intercepts" the functions' arguments and converts a certain keyword argument from one type to another before passing it on.

Here's a minimal example of what I'm working with, with partial annotations:

from __future__ import annotations
from functools import wraps


class User:
    pass


class Result:
    pass


def get_user_from_database(username: str) -> User | None:
    raise NotImplementedError


def fetch_user(f):
    @wraps(f)
    def wrapped(*, username: str, **kwargs) -> Result:
        if (user := get_user_from_database(username)) is not None:
            return f(user=user, **kwargs)
        else:
            return Result()

    return wrapped


@fetch_user
def function1(*, user: User) -> Result:
    raise NotImplementedError

# function1() can now be called with a `username: str` argument,
# and the decorator will pass a `user: User` argument to the function body.

@fetch_user
def function2(*, user: User, flavor: str) -> Result:
    raise NotImplementedError

Here's my first attempt at annotating this, using Concatenate and ParamSpec:

from __future__ import annotations
from collections.abc import Callable
from functools import wraps
from typing import Concatenate, ParamSpec


class User:
    pass


class Result:
    pass


def get_user_from_database(username: str) -> User | None:
    raise NotImplementedError


P = ParamSpec("P")


def fetch_user(
    f: Callable[Concatenate[User, P], Result]
) -> Callable[Concatenate[str, P], Result]:
    @wraps(f)
    def wrapped(*, username: str, **kwargs: P.kwargs) -> Result:
        if (user := get_user_from_database(username)) is not None:
            return f(user=user, **kwargs)
        else:
            return Result()

    return wrapped


@fetch_user
def function1(*, user: User) -> Result:
    raise NotImplementedError


@fetch_user
def function2(*, user: User, flavor: str) -> Result:
    raise NotImplementedError

Running mypy 1.6.0 with no configuration on this code, using Python 3.11.6 on macOS 14.0, gives me the following errors:

attempt00.py:28: error: Too few arguments  [call-arg]
attempt00.py:28: error: Argument "user" has incompatible type "User"; expected "P.kwargs"  [arg-type]
attempt00.py:32: error: Incompatible return value type (got "Callable[[NamedArg(str, 'username'), KwArg(Any)], Result]", expected "Callable[[str, **P], Result]")  [return-value]
attempt00.py:35: error: Argument 1 to "fetch_user" has incompatible type "Callable[[NamedArg(User, 'user')], Result]"; expected "Callable[[User, NamedArg(User, 'user')], Result]"  [arg-type]
attempt00.py:40: error: Argument 1 to "fetch_user" has incompatible type "Callable[[NamedArg(User, 'user'), NamedArg(str, 'flavor')], Result]"; expected "Callable[[User, NamedArg(User, 'user'), NamedArg(str, 'flavor')], Result]"  [arg-type]
Found 5 errors in 1 file (checked 1 source file)

---

Here's my second attempt at annotations, this time using Procotols with __call__ methods that take **kwargs: Unpack[Kwargs]:

from __future__ import annotations
from functools import wraps
from typing import Generic, Protocol, TypeVar, TypedDict, Unpack


class User:
    pass


class Result:
    pass


def get_user_from_database(username: str) -> User | None:
    raise NotImplementedError


Kwargs = TypeVar("Kwargs", bound=TypedDict, covariant=True)


class UsernameFunc(Protocol, Generic[Kwargs]):
    def __call__(self, *, username: str, **kwargs: Unpack[Kwargs]) -> Result:
        ...


class UserFunc(Protocol, Generic[Kwargs]):
    def __call__(self, *, user: User, **kwargs: Unpack[Kwargs]) -> Result:
        ...


def fetch_user(f: UsernameFunc[Kwargs]) -> UserFunc[Kwargs]:
    @wraps(f)
    def wrapped(*, username: str, **kwargs: Unpack[Kwargs]) -> Result:
        if (user := get_user_from_database(username)) is not None:
            return f(user=user, **kwargs)
        else:
            return Result()

    return wrapped


@fetch_user
def function1(*, user: User) -> Result:
    raise NotImplementedError


@fetch_user
def function2(*, user: User, flavor: str) -> Result:
    raise NotImplementedError

Running mypy --enable-incomplete-feature=Unpack on this code gives me the following errors:

attempt01.py:18: error: Variable "typing.TypedDict" is not valid as a type  [valid-type]
attempt01.py:18: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
attempt01.py:22: error: Unpack item in ** argument must be a TypedDict  [misc]
attempt01.py:27: error: Unpack item in ** argument must be a TypedDict  [misc]
attempt01.py:33: error: Unpack item in ** argument must be a TypedDict  [misc]
attempt01.py:42: error: Argument 1 to "fetch_user" has incompatible type "Callable[[NamedArg(User, 'user')], Result]"; expected "UserFunc[<nothing>]"  [arg-type]
attempt01.py:47: error: Argument 1 to "fetch_user" has incompatible type "Callable[[NamedArg(User, 'user'), NamedArg(str, 'flavor')], Result]"; expected "UserFunc[<nothing>]"  [arg-type]
Found 6 errors in 1 file (checked 1 source file)

I also tried a variant of this latter attempt where the bound on Kwargs is changed to the class class KwargsClass(TypedDict): pass, which has a subclass that adds a flavor: str attribute, but the only difference that made was getting rid of the Variable "typing.TypedDict" is not valid as a type error.

So, is there a way to properly annotate this, and what is it?

[erictraut]

I don't think there's a way to annotate this using today's type system. It's possible to replace the type of a positional parameter using a ParamSpec and Concatenate, but this technique doesn't work for keyword parameters.

Your first attempt doesn't work because **kwargs: P.kwargs must always be used in conjunction with **args: P.args. The reasons for this are spelled out in PEP 612. Mypy has a bug where it does not report an error for this condition.

Your second attempt doesn't work for a couple of reasons. First, you've defined a TypeVar with an upper bound of TypedDict. This doesn't make sense because TypedDict is a special form in the Python type system that is used to define a structural type. It doesn't participate in the type system itself, so a class created with TypedDict doesn't "inherit from" TypedDict. Second, PEP 692 clearly specifies that **kwargs can be used only with an unpacked TypedDict. You cannot use it with a TypeVar or other type forms (like unions, etc.).