Python 3 type hinting for decorator

2020-02-10 12:23发布

问题:

Considere the following code:

from typing import Callable, Any

TFunc = Callable[..., Any]

def get_authenticated_user(): return "John"

def require_auth() -> Callable[TFunc, TFunc]:
    def decorator(func: TFunc) -> TFunc:
        def wrapper(*args, **kwargs) -> Any:
            user = get_authenticated_user()
            if user is None:
                raise Exception("Don't!")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@require_auth()
def foo(a: int) -> bool:
    return bool(a % 2)

foo(2)      # Type check OK
foo("no!")  # Type check failing as intended

This piece of code is working as intended. Now imagine I want to extend this, and instead of just executing func(*args, **kwargs) I want to inject the username in the arguments. Therefore, I modify the function signature.

from typing import Callable, Any

TFunc = Callable[..., Any]

def get_authenticated_user(): return "John"

def inject_user() -> Callable[TFunc, TFunc]:
    def decorator(func: TFunc) -> TFunc:
        def wrapper(*args, **kwargs) -> Any:
            user = get_authenticated_user()
            if user is None:
                raise Exception("Don't!")
            return func(*args, user, **kwargs)  # <- call signature modified

        return wrapper

    return decorator


@inject_user()
def foo(a: int, username: str) -> bool:
    print(username)
    return bool(a % 2)


foo(2)      # Type check OK
foo("no!")  # Type check OK <---- UNEXPECTED

I can't figure out a correct way to type this. I know that on this example, decorated function and returned function should technically have the same signature (but even that is not detected).

回答1:

You can't use Callable to say anything about additional arguments; they are not generic. Your only option is to say that your decorator takes a Callable and that a different Callable is returned.

In your case you can nail down the return type with a typevar:

RT = TypeVar('RT')  # return type

def inject_user() -> Callable[[Callable[..., RT]], Callable[..., RT]]:
    def decorator(func: Callable[..., RT]) -> Callable[..., RT]:
        def wrapper(*args, **kwargs) -> RT:
            # ...

Even then the resulting decorated foo() function has a typing signature of def (*Any, **Any) -> builtins.bool* when you use reveal_type().

Various proposals are currently being discussed to make Callable more flexible but those have not yet come to fruition. See

  • Allow variadic generics
  • Proposal: Generalize Callable to be able to specify argument names and kinds
  • TypeVar to represent a Callable's arguments
  • Support function decorators excellently

for some examples. The last one in that list is an umbrella ticket that includes your specific usecase, the decorator that alters the callable signature:

Mess with the return type or with arguments

For an arbitrary function you can't do this at all yet -- there isn't even a syntax. Here's me making up some syntax for it.