How to properly annotate a decorator that decorates both regular functions and async functions?

[MapleCCC]

I am writing a decorator that caches the result of the decorated function. Should be a trivial use case. The decorator supports both regular function and async function. However, I don't seem to be able to find a way to add type annotations that can make Pyright happy.

My original code was like this:

import inspect
from collections.abc import Callable
from functools import wraps
from typing import TypeVar

from typing_extensions import ParamSpec

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

def cache(func: Callable[P, R]) -> Callable[P, R]:

    if inspect.iscoroutinefunction(func):

        @wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            result = await func(*args, **kwargs)
            # Omitted the code to add the result to cache
            return result

    else:

        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            result = func(*args, **kwargs)
            # Ommitted the code to add the result to cache
            return result

    return wrapper

Pyright complained that the the type R is not awaitable, so I make some changes to comply:

import inspect
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import TypeVar, Union

from typing_extensions import ParamSpec

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


def cache(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> Union[Callable[P, R], Callable[P, Awaitable[R]]]:

    if inspect.iscoroutinefunction(func):

        @wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            result = await func(*args, **kwargs)
            # Omitted the code to add the result to cache
            return result

    else:

        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            result = func(*args, **kwargs)
            # Ommitted the code to add the result to cache
            return result

    return wrapper

Pyright still complained that the type R is not awaitable. It seems the problem originates from that Pyright didn't narrow the type of func after the inspect.iscoroutinefunction call?

I don't know how to fix the type annotations here. I find adding type annotations in situations involving higher order functions quite tricky.

[erictraut]

One problem that I see here is that the union includes Callable[P, R] and Callable[P, Awaitable[R]], but since R is unbound and unconstrained, an Awaitable return type will match both subtypes in the union. In other words, the constraint solver will be allowed to select either subtype, and the result will be nondeterministic (based on the internal heuristics or implementation details of the constraint solver). Pyright's heuristics always attempt to choose the "more complex" subtype when there's ambiguity, but this behavior is not dictated by PEP 484 or any other typing standard.

Another problem is that pyright doesn't support iscoroutinefunction for narrowing. The current list of supported type guard patterns is documented here. We do occasionally add new ones, but each new pattern requires specific logic that adds complexity and potentially slows down analysis. For those reasons, we add new type guard patterns only when we have a strong signal that it's a commonly-used pattern. That can take the form of upvotes or support in other type checkers like mypy. I haven't seen this pattern in any code before, and mypy doesn't support it, so it's probably not commonly used. We could add support for it in the future if we get enough signal from pyright users.

Even if pyright supported type narrowing based on iscoroutinefunction, there's another issue that would prevent it from working in the code above. Type narrowing is not applied to captured variables that are referenced from within inner scopes. In your code sample, the variable func is captured by the inner scope of wrapper. The type checker has no way of knowing when wrapper will be executed, so it must assume that the value of the captured value could be modified prior to execution. Because of this, it always uses the unnarrowed type within the captured scope. You can work around this by assigning the narrowed value to another local variable that is referenced from within the inner scope.

Here are a couple of workarounds for you to consider.

First, you can simply use a cast.

awaitable_func = cast(Callable[P, Awaitable[R]], func)
result = await awaitable_func(*args, **kwargs)

Second, you can use a user-defined type guard (PEP 647).

def is_coroutine(func: Union[Callable[P, R], Callable[P, Awaitable[R]]]) -> TypeGuard[Callable[P, Awaitable[R]]]:
    return inspect.iscoroutinefunction(func)

def cache(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> Union[Callable[P, R], Callable[P, Awaitable[R]]]:

    if is_coroutine(func):
        # Narrowed value assigned to new local variable
        awaitable_func = func

        @wraps(awaitable_func)
        async def wrapper1(*args: P.args, **kwargs: P.kwargs) -> R:
            result = await awaitable_func(*args, **kwargs)
            # Omitted the code to add the result to cache
            return result

        return wrapper1
    else:

        @wraps(func)
        def wrapper2(*args: P.args, **kwargs: P.kwargs) -> R:
            result = func(*args, **kwargs)
            # Ommitted the code to add the result to cache
        return wrapper2

[MapleCCC]

@erictraut Very thank you for your patient and informative reply. It helps me understand the concepts a lot. Didn't know about that type narrowing is not applied to captured variable referenced in inner scopes. Quite an advanced knowledge to me.

I tried the user-defined type guard approach and the cast approach per your suggestions, they both work perfectly. It's great to finally manage to make this work. This type annotation thing is not very beginner-friendly I guess 🤣.

Now that the problem is resolved. I will close the issue then. Again thank you for your kind help and patient reply ❤️.

[zanieb]

Hi! I'm also trying to write an async/sync decorator and the type narrowing is not working as described. The return types remain a Union.

Here's a complete example created from the samples above

import inspect
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import TypeVar, Union

from typing_extensions import ParamSpec, TypeGuard

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


def is_coroutine(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> TypeGuard[Callable[P, Awaitable[R]]]:
    return inspect.iscoroutinefunction(func)


def cache(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> Union[Callable[P, R], Callable[P, Awaitable[R]]]:

    if is_coroutine(func):
        # Narrowed value assigned to new local variable
        awaitable_func = func
        reveal_type(awaitable_func)

        @wraps(awaitable_func)
        async def wrapper1(*args: P.args, **kwargs: P.kwargs) -> R:
            result = await awaitable_func(*args, **kwargs)
            # Omitted the code to add the result to cache
            return result

        reveal_type(wrapper1)
        return wrapper1
    else:
        reveal_type(func)

        @wraps(func)
        def wrapper2(*args: P.args, **kwargs: P.kwargs) -> R:
            result = func(*args, **kwargs)
            # Ommitted the code to add the result to cache
            return result

        reveal_type(wrapper2)
        return wrapper2


@cache
def foo(x: int) -> str:
    return str(x)


@cache
async def bar(x: int) -> str:
    return str(x)


f = foo(1)
b = bar(2)
reveal_type(f)
reveal_type(b)

Output

❯ pyright example.py
No configuration file found.
No pyproject.toml file found.
stubPath ./typings is not a valid directory.
Assuming Python platform Darwin
Searching for source files
Found 1 source file
./example.py
  ./example.py:25:21 - info: Type of "awaitable_func" is "() -> Awaitable[R@cache]"
  ./example.py:33:21 - info: Type of "wrapper1" is "(*args: P.args.args, **kwargs: P.kwargs.kwargs) -> Coroutine[Any, Any, R@cache]"
  ./example.py:36:21 - info: Type of "func" is "Callable[P@cache, R@cache] | Callable[P@cache, Awaitable[R@cache]]"
  ./example.py:42:20 - error: Expression of type "R@cache | Awaitable[R@cache]" cannot be assigned to return type "R@cache"
    Type "R@cache | Awaitable[R@cache]" cannot be assigned to type "R@cache" (reportGeneralTypeIssues)
  ./example.py:44:21 - info: Type of "wrapper2" is "(*args: P.args.args, **kwargs: P.kwargs.kwargs) -> R@cache"
  ./example.py:60:13 - info: Type of "f" is "str | Awaitable[str]"
  ./example.py:61:13 - info: Type of "b" is "Coroutine[Any, Any, str] | Awaitable[Coroutine[Any, Any, str]]"
1 error, 0 warnings, 6 infos 
Completed in 0.526sec

We can add a cast(R, result) to the return in wrapper2 to resolve the error, but the narrowing still does not apply. Additionally, adding cast to the returned wrappers (e.g. return cast(Callable[P, Awaitable[R]], wrapper1)) does not make a difference.

Versions

❯ pyright --version 
pyright 1.1.165
❯ python --version           
Python 3.9.6

[MapleCCC]

@madkinsz Thanks for the follow-up. I was the original OP. At the time I tried several workarounds, and despite that they worked, the whole process of wrestling with type system is tedious and lengthy. It's surprisingly difficult to write proper and correct type annotations for some trivial use cases.

Even if you manage to eliminate all type errors, the source code could have by then inadvertently become too mouthful and about 3x longer and complicated. At the end, I decided to trade type annotation completeness for simplicity.

My recommendation is to use several @overload-s to enumerate all the possible function signatures, while dropping type annotations in the actual implementation. This way the code looks cleaner and easier to read.

It's more promising and productive to wait for the Python dev team to, in the future, roll out more sophisticated type system constructs that can express more powerful type concepts, and hopefully make our case here easier to annotate. And before that happens, I suggest we stick with the simplicity-first philosophy, and refrain from pursuing perfect type annotation completeness.

[erictraut]

@madkinsz, are you expecting type narrowing to be applied with the expression foo(1)? Type narrowing doesn't apply in this case. It applies only in the cases of assignment expressions and conditional checks (like if and while conditions). This is a call expression, so the signature of the call dictates the result type. The cache decorator returns two function signatures, and they indicate return types R and Awaitable[R], respectively. If you call this union of functions, the result will be a union of their return types. That's why you're seeing a union here.

As @MapleCCC mentioned, the proper solution here is to define two overload functions for cache.

@overload
def cache(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
    ...

@overload
def cache(func: Callable[P, R]) -> Callable[P, R]:
    ...

[zanieb]

Thanks for the quick reply! It's definitely tricky working with dual sync/async code, and even more so once you introduce type checking. I've been close to giving up a few times now, but since I'm working on a library that will be consumed by lots of beginners, any work I can do to provide proper hinting in an IDE is worth the complexity.

I was hoping the type narrowing would apply within the decorator and that it would carry to the foo(1) expression as an inferred return type. Thanks for clarifying that it's not expected to do so.

I've been working with overload a lot and I'm running into a separate issue there with Any / untyped functions. If there's an Any return type we match the first overload and mark it as async instead of sync.

@overload
def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[List[R]]]:
    ...


@overload
def decorator(func: Callable[P, R]) -> Callable[P, List[R]]:
    ...


def decorator(func):
    if inspect.iscoroutinefunction(func):

        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> List[R]:
            return [func(*args, **kwargs)]

        return cast(Callable[P, Awaitable[List[R]]], async_wrapper)

    else:

        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> List[R]:
            return [func(*args, **kwargs)]

        return cast(Callable[P, List[R]], sync_wrapper)


@decorator
def sync_fn(x: int) -> str:
    return str(x)


@decorator
async def async_fn(x: int) -> str:
    return str(x)


sync_result = sync_fn(1)
async_result = async_fn(2)
reveal_type(sync_result)  # List[str]
reveal_type(async_result)  # Awaitable[List[str]]


@decorator
def sync_fn_any(x):
    return x


@decorator
async def async_fn_any(x):
    return x


sync_fn_any_result = sync_fn_any(1)
async_fn_any_result = async_fn_any(1)
reveal_type(sync_fn_any_result)    # Awaitable[List[Unknown]]    -- Wanted List[Any] or List[Unknown]
reveal_type(async_fn_any_result)  # Awaitable[List[Unknown]]

Adding another overload covering Any explicitly does not work (and I did not really expect it to)

@overload
def decorator(func: Callable[P, Any]) -> Callable[P, List[Any]]:
    ...

[erictraut]

Yeah, overloads can be a problem if the input arguments are untyped. You could add another overload with a NoReturn type. The only type that matches a NoReturn is Any, so it will take care of the Any case.

@overload
def decorator(func: Callable[P, NoReturn]) -> Callable[P, list]:
    ...

You can also eliminate those unsightly cast calls as follows:

def decorator(
    func: Callable[P, R]
) -> Union[Callable[P, List[R]], Callable[P, Awaitable[List[R]]]]:
    if is_coroutine(func):

        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> List[R]:
            return [func(*args, **kwargs)]

        return async_wrapper

    else:

        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> List[R]:
            return [func(*args, **kwargs)]

        return sync_wrapper

[zanieb]

Wonderful! Thank you for the tip on NoReturn -- that fixes a big issue for me.

I still need to include the casts, for some reason the parameters are () otherwise and I get an error. Perhaps this is a regression like you mentioned in #2268? Regardless, my real use case is a bit more complicated as I have a decorator that also takes settings and I've got to have a few casts to make it work anyway.

[maresb]

This thread has been extremely enlightening. To echo @madkinsz, it looks like a cast is still necessary for the sync case. The reason seems to be that the TypeGuard applies only when is_coroutine is true, so in the sync case, func is not narrowed, and its return value is a priori Union[R, Awaitable[R]]. This is a problem since sync_wrapper is not allowed to return Awaitable[R].

Here is complete code illustrating the one necessary cast:

import inspect
from typing import (
    Awaitable,
    Callable,
    ParamSpec,
    TypeGuard,
    TypeVar,
    Union,
    cast,
    overload,
)

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


def is_coroutine(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> TypeGuard[Callable[P, Awaitable[R]]]:
    return inspect.iscoroutinefunction(func)


@overload
def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
    ...


@overload
def decorator(func: Callable[P, R]) -> Callable[P, R]:
    ...


def decorator(
    func: Union[Callable[P, R], Callable[P, Awaitable[R]]]
) -> Union[Callable[P, R], Callable[P, Awaitable[R]]]:
    if is_coroutine(func):

        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            return await func(*args, **kwargs)

        return async_wrapper
    else:
        func = cast(Callable[P, R], func)

        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            return func(*args, **kwargs)

        return sync_wrapper

[zanieb]

Hey @askerka have you reported that as an issue to mypy? I'm not sure this is the right place to raise that issue.

[askerka]

@madkinsz Sorry, I've overlooked the thread place and already deleted the comment