Need help with typing a decorator that can work on both regular functions and methods.

[vamshiaruru-virgodesigns]

Hi all, I am trying to type a decorator that adds an additional functional argument using ParamSpec, and Concatenate. I am able to get it working for regular functions, but that same type hint doesn't work for method. I can get it working for a class method as well, but in that case it won't work for regular functions. Is there a way to get it working for both? Example below:

import functools
from typing import Union, overload, Callable, TypeVar, Optional, Union
from typing_extensions import Concatenate, ParamSpec

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

class Client:
   ...

Regular = Callable[[Callable[Concatenate[Client, P], R]], Callable[P, R]]
Method = Callable[[Callable[Concatenate[Self, Client, P], R]], Callable[Concatenate[Self, P], R]]


@overload
def inject_client() -> Method:
    ...

@overload
def inject_client(url: Optional[str], auth_token: Optional[str]) -> Method:
    ...

@overload
def inject_client() -> Regular:
    ...

@overload
def inject_client(url: Optional[str], auth_token: Optional[str]) -> Regular:
    ...
 
 
 def inject_client(
    url=None,
    auth_token=None
) -> Union[Method, Regular]:
    def wrap(function):
        @functools.wraps(function)
        async def wrapper(*args, **kwargs):
            async with Client(url, auth_token) as client:
                return await function(client=client, *args, **kwargs)

        return wrapper

    return wrap # type: ignore

The above works when used on a method like so:

class Test:
   @inject_client
   def test(self, client: Client):
     ...

But when I use it on regular functions like this:

@inject_client()
def testing(client: Client):
   pass

I get the error from pyright:

Argument of type "(client: Client) -> None" cannot be assigned to parameter of type "(Unknown, Client, ...) -> Unknown"
  Type "(client: Client) -> None" cannot be assigned to type "(Unknown, Client, ...) -> Unknown"
    Function accepts too many positional parameters; expected 1 but received 2

Any help is appreciated. TIA :)

[Drino]

Hello!

I'm sorry about you having to wait for answer this long, hope it still will be useful.

There are two issues in the snippet:

1. Callable type alias doesn't work as expected (to shortcut decorator type it's better to define protocol)
2. @overload does dispatch by the argument list (so @overload argument list should be different - in the snippet the only difference is the return type)

So, to fix that we need:

1. Rewrite Regular and Method as Protocols.
2. Use @overload on different argument lists. Actually, the right place is the __call__ of this Protocol.

So it will be something like:

import functools
from typing import Union, overload, Callable, TypeVar, Optional, Union, Protocol
from typing_extensions import Concatenate, ParamSpec

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

class Client:
   ...


class ClientInjector(Protocol):
    @overload
    def __call__(self, _: Callable[Concatenate[Self, Client, P], R], /) -> Callable[Concatenate[Self, P], R]:
        ...
    
    @overload
    def __call__(self, _: Callable[Concatenate[Client, P], R], /) -> Callable[P, R]:
        ...
    
    def __call__(self, _: Callable) -> Callable:
        ...


def inject_client(
    url: Optional[str] =None,
    auth_token: Optional[str] = None
) -> ClientInjector:
    def wrap(function):
        @functools.wraps(function)
        async def wrapper(*args, **kwargs):
            async with Client(url, auth_token) as client:
                return await function(client=client, *args, **kwargs)

        return wrapper

    return wrap # type: ignore
    
    
class Test:
   @inject_client()
   def test(self, client: Client):
     ...


@inject_client()
def testing(client: Client):
   pass

reveal_type(Test.test)
reveal_type(testing)

https://mypy-play.net/?mypy=latest&python=3.10&gist=bed0501ec912ca17a6bad96c0d49bfd3