How to type annotate a multi-level decorator

Question

I'm trying to annotate an injector decorator that injects a value from a global dictionary as a keyword argument into the decorated function when the function is called.

Can anyone experienced with annotating decorators with parameters help me out? Tried annotating but got stuck on the errors below:

import functools
import inspect
from typing import Any, Callable, TypeVar, ParamSpec


Type = TypeVar('Type')
Param = ParamSpec('Param')
_INSTANCES = {}


def make_injectable(instance_name: str, instance: object) -> None:
    _INSTANCES[instance_name] = instance


def inject(*instances: str) -> Callable[Param, Type]:

    def get_function_with_instances(fn: Callable[Param, Type]) -> Callable[Param, Type]:
        # This attribute is to easily access which arguments of fn are injectable
        fn._injectable_args = instances

        def handler(*args: Param.args, **kwargs: Param.kwargs) -> Type:
            new_kwargs: dict[str, Any] = dict(kwargs).copy()
            for instance in instances:
                if instance in new_kwargs:
                    continue
                if instance not in _INSTANCES:
                    raise ValueError(f"Instance {instance} was not initialized yet")
                new_kwargs[instance] = _INSTANCES[instance]
            return fn(*args, **new_kwargs)

        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> Callable[Param, Type]:
                return await handler(*args, **kwargs)

        else:
            @functools.wraps(fn)
            def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> Callable[Param, Type]:
                return handler(*args, **kwargs)

        return wrapper

    return get_function_with_instances

If I run mypy with these annotations I get these errors I cannot circumvent without creating new ones:

mypy injector.py --strict --warn-unreachable --allow-subclassing-any --ignore-missing-imports --show-error-codes --install-types --non-interactive

injector.py:33: error: "Callable[Param, Type]" has no attribute "_injectable_args"  [attr-defined]
injector.py:48: error: Returning Any from function declared to return "Callable[Param, Type]"  [no-any-return]
injector.py:48: error: Incompatible types in "await" (actual type "Type", expected type "Awaitable[Any]")  [misc]
injector.py:53: error: Incompatible return value type (got "Type", expected "Callable[Param, Type]")  [return-value]
injector.py:55: error: Incompatible return value type (got "Callable[Param, Coroutine[Any, Any, Callable[Param, Type]]]", expected "Callable[Param, Type]")  [return-value]
injector.py:57: error: Incompatible return value type (got "Callable[[Callable[Param, Type]], Callable[Param, Type]]", expected "Callable[Param, Type]")  [return-value]

Thank you for your time.

Answer 1

The first [attr-defined] error is unavoidable IMO and should be simply explicitly ignored. There is no point in reinventing the wheel and defining your own special callable protocol with that special attribute.

The second and third errors with the codes [no-any-return]/misc I'll come back to later.

That fourth error with the [return-value] code comes up because the return annotation of your wrapper should be T, not Callable[P, T]. It is supposed to return whatever the decorated function returns after all.

The fifth error (also [return-value]) tells you that wrapper may be a coroutine that can be awaited to yield T, but you declared get_function_with_instances to return a callable that returns T (not a coroutine to await T from).

The very last [return-value] error comes up because inject returns a decorator that takes an argument of type Callable[P, T] and returns an object of that same type again. So the return annotation for inject should indeed be Callable[[Callable[P, T]], Callable[P, T]], just as mypy says.

---

Now for the [no-any-return]/misc errors. This gets a bit confusing since your intention was to cover both the case of fn being a coroutine function and the case of it being a regular function.

You annotate handler to return T just like fn. But what that T is, is not further narrowed. The type guard given by iscoroutinefunction applies to fn and does not automatically extend to handler. From the point of view of the static type checker, handler returns some object. And that cannot be safely assumed to be awaitable. Therefore you cannot safely use it with await (the [misc] error).

Since the type checker doesn't even allow the await expression in that line, it obviously cannot verify that the returned value actually matches the annotation of wrapper (which should be T just like with the error mentioned earlier, but in this case it doesn't matter either way).

I am not 100 % sure about the root cause those two errors though.

---

If I were you, I would make my life a whole lot easier by not even inspecting the decorated function in the first place. The behavior of your decorator does not change. The only difference is that one call needs to be followed by an await to get the value. You can let the decorator be agnostic as to whether fn returns an awaitable or not and leave it up to the caller to handle it.

Here is my suggestion:

from collections.abc import Callable
from functools import wraps
from typing import TypeVar, ParamSpec


T = TypeVar('T')
P = ParamSpec('P')
_INSTANCES = {}


def make_injectable(instance_name: str, instance: object) -> None:
    _INSTANCES[instance_name] = instance


def inject(*instances: str) -> Callable[[Callable[P, T]], Callable[P, T]]:
    def get_function_with_instances(fn: Callable[P, T]) -> Callable[P, T]:
        # This attribute is to easily access which arguments of fn are injectable
        fn._injectable_args = instances  # type: ignore[attr-defined]

        @wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            for instance in instances:
                if instance in kwargs:
                    continue
                if instance not in _INSTANCES:
                    raise ValueError(f"Instance {instance} was not initialized yet")
                kwargs[instance] = _INSTANCES[instance]
            return fn(*args, **kwargs)
        return wrapper
    return get_function_with_instances

Here is quick test to show that the types are all correctly inferred:

make_injectable("foo", object())


@inject("foo")
def f(**kwargs: object) -> int:
    print(kwargs)
    return 1


@inject("foo")
async def g(**kwargs: object) -> int:
    print(kwargs)
    return 2


async def main() -> tuple[int, int]:
    x = f()
    y = await g()
    return x, y


if __name__ == '__main__':
    from asyncio import run
    print(run(main()))

Sample output:

{'foo': <object object at 0x7fe39fea0b20>}
{'foo': <object object at 0x7fe39fea0b20>}
(1, 2)

There are no complaints from mypy --strict. The way main is written, we can see that the return types are all inferred correctly, but if we wanted to check explicitly, we could add reveal_type(f) and reveal_type(g) at the end of the script. Then mypy would tell us:

Revealed type is "def (**kwargs: builtins.object) -> builtins.int"
Revealed type is "def (**kwargs: builtins.object) -> typing.Coroutine[Any, Any, builtins.int]"

This is what we expected.