CODEWITHSUNDEEP
pythontypestype-hintingmypypep
rednafi
rednafi

Reputation: 1731

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

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.

Upvotes: 5

Views: 2826

Answers (1)

Mario Ishac
Mario Ishac

Reputation: 5907

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.

Upvotes: 1

Related Questions