Typing a Decorator that Accepts Arguments via ParamSpec (PEP-612, Python 3.10)

Question

I was reading PEP-612 and it makes typing a decorator fairly tractable. Also, the example provided in the PEP makes it look pretty easy. This example is copied directly from the PEP:

from typing import ParamSpec, TypeVar
from collections.abc import Callable, Awaitable

P = ParamSpec("P")
R = TypeVar("R")

def add_logging(f: Callable[P, R]) -> Callable[P, Awaitable[R]]:
  async def inner(*args: P.args, **kwargs: P.kwargs) -> R:
    await log_to_database()
    return f(*args, **kwargs)
  return inner

@add_logging
def takes_int_str(x: int, y: str) -> int:
  return x + 7

await takes_int_str(1, "A") # Accepted
await takes_int_str("B", 2) # Correctly rejected by the type checker

However, I found it non-trivial to properly type annotate a parametrizable decorator. Check the following MRE:

import functools
import inspect 
from collections.abc import Callable, Coroutine
from typing import ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

def tag(*names: str) -> ??:
    """This decorator just tags an async function with the provided name(s)."""

    for name in names:
        if not isinstance(name, str):
            raise TypeError("tag name must be a string")

    def outer(func: Callable[P, R]) -> Callable[P, Coroutine[R]]:
        func.tags = names
        if not inspect.iscoroutinefunction(func):
            raise TypeError("tagged function must be an async function")

        @functools.wraps(func)
        async def inner(*args: P.args, **kwargs: P.kwargs) -> R:
            result = await func(*args, **kwargs)
            return result

        return inner

    return outer

I'm struggling with figuring out the return type of the tag function. Also, I'm not 100% confident with the correctness of the typing of outer and inner nested functions. How do I type this properly?

P.S. I know, as of today, mypy 0.902 doesn't support this feature fully yet.

a) I believe `func`'s type has to be `Callable[P, Coroutine[R]]`(note the `Coroutine`) because you `await` its result in `inner`. An `R`, on the other hand, doesn't guarantee that it's awaitable. b) I think `tag`'s return type is `Callable[[Callable[P, Coroutine[R]]], Callable[P, Coroutine[R]]]` because `tag` returns a decorator that takes a coroutine function and returns a coroutine function of the same signature. — Mario Ishac, Jun 17 '21 at 21:45
Yes. I think you're correct about the `func`'s type. I guess I was missing the list inside `tag`'s Callable. However, there's no way of being sure until *mypy* fully supports this. Nonetheless, do you mind adding the comment as an answer for documentation purposes? Thanks you. — Redowan Delowar, Jun 17 '21 at 21:51

score 1 · Accepted Answer · answered Jun 17 '21 at 22:08

First thing to note is that your parametrized decorator example isn't just the PEP's decorator example plus parametrization. Instead, your second example's decorator (after parametrization) takes a asynchronous function, whereas the PEP's example takes a synchronous function.

Because you are awaiting the func's result directly, unlike the PEP's example which awaited a separate logger followed by calling f normally, your outer needs to take a Callable[[P], Awaitable[R]] instead of Callable[[P], R].

Second thing to note is that regarding tags return type, you can figure it out by adding reveal_type(outer), which would in turn be the return type of tag. I haven't run this (because of mypy not actually supporting your example yet), but it should say Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]. In other words, tag returns a decorator that itself takes an asynchronous function and returns an asynchronous function.

Third thing to note is that you probably want Awaitable[T] in all your examples (which is why I've been using it myself throughout the above) instead of Coroutine[T]. This is because a) Coroutine takes three type parameters instead of one (so you'd have to use Coroutine[Any, Any, T] instead of Coroutine[T], where the first two type parameters are for send) and b) Coroutine is a subtype of Awaitable with the added support of sending, which you don't utilize anyway.

The `Coroutine` vs `Awaitable` comparison was fantastic. I thought I was narrowing down the type more by using `Coroutine` instead of `Awaitable`. However, now I'm convinced. Thanks a ton. Also, I made the second example intentionally different from the first one to see how tackle this sort of gnarly situation. Your explanation makes a lot of sense. — Redowan Delowar, Jun 17 '21 at 22:20

Typing a Decorator that Accepts Arguments via ParamSpec (PEP-612, Python 3.10)

1 Answers1