Properly type hinting a decorator which creates async context managers

[ronf]

I've written a decorator which accepts a coroutine function that returns a class which supports being an async context manager and lets me directly use "async with" on it, without having to manually await it first. So, instead of writing:

    conn = await connect(host)

    async with conn:
        # do something here with conn.

I can directly write:

    async with connect(host) as conn:
        # do something here with conn.

The type returns by awaiting connect() is an async context manager (with __aenter__ and __aexit__). The decorator just passes through the calls to those methods when wrapping the coroutine.

In fact, the decorator allows either syntax to be used on the wrapped coroutine. This is working great, but I noticed that when I feed this through mypy, it complains that it can't find __aenter__ or __aexit__ on the returned value. I think this is part because I tried to preserve the original function signature for typing purposes in the decorator. I've tried using generics to more accurately specify the types and the relationship between the input function (a coroutine function) and the output (an async context manager of the coroutine function's return value), but so far I haven't gotten it right. Here's an example of the errors mypy returns:

callback_client2.py:34: error: "Coroutine[Any, Any, SSHClientConnection]" has no attribute "__aenter__"
callback_client2.py:34: error: "Coroutine[Any, Any, SSHClientConnection]" has no attribute "__aexit__"

Here's the code I've been using with same very rudimentary type hints which at leasts lets it type check the wrapped functions as if they weren't decorated. Any thoughts on getting rid of the errors above, though? Is there some way I can easily have the decorator add the necessary annotations for the new methods?

DecoratedFunc = TypeVar('DecoratedFunc', bound=Callable)

def async_context_manager(coro: DecoratedFunc) -> DecoratedFunc:
    """Decorator for functions returning asynchronous context managers

       This decorator can be used on functions which return objects
       intended to be async context managers. The object returned by
       the function should implement __aenter__ and __aexit__ methods
       to run when the async context is entered and exited.

       This wrapper also allows the use of "await" on the function being
       decorated, to return the context manager without entering it.

    """

    class AsyncContextManager:
        """Async context manager wrapper"""

        def __init__(self, coro):
            self._coro = coro
            self._result = None

        def __await__(self):
            return self._coro.__await__()

        async def __aenter__(self):
            self._result = await self._coro
            return await self._result.__aenter__()

        async def __aexit__(self, *exc_info):
            await self._result.__aexit__(*exc_info)
            self._result = None
            return False

    @functools.wraps(coro)
    def context_wrapper(*args, **kwargs):
        """Return an async context manager wrapper for this coroutine"""

        return AsyncContextManager(coro(*args, **kwargs))

    return cast(DecoratedFunc, context_wrapper)

[ronf]

I've made some progress on this. The latest version I have is:

import asyncio
import functools
from types import TracebackType
from typing import Any, AsyncContextManager, Awaitable, Callable, Generator
from typing import Generic, Optional, Type, TypeVar, cast

_ACM = TypeVar('_ACM', bound=AsyncContextManager)

class _ACMWrapper(Generic[_ACM]):
    """Async context manager wrapper"""

    def __init__(self, coro: Awaitable[_ACM]):
        self._coro = coro
        self._coro_result: Optional[_ACM] = None

    def __await__(self) -> Generator[Any, None, _ACM]:
        return self._coro.__await__()

    async def __aenter__(self) -> _ACM:
        self._coro_result = await self._coro

        return await self._coro_result.__aenter__()

    async def __aexit__(self, exc_type: Type[BaseException],
                        exc_value: BaseException,
                        traceback: TracebackType) -> bool:
        exit_result = await self._coro_result.__aexit__(
            exc_type, exc_value, traceback)

        self._coro_result = None

        return exit_result

_ACMCoroFunc = Callable[..., Awaitable[_ACM]]
_ACMWrapperFunc = Callable[..., _ACMWrapper[_ACM]]

def async_context_manager(coro: _ACMCoroFunc[_ACM]) -> _ACMWrapperFunc[_ACM]:
    """Decorator for functions returning asynchronous context managers

       This decorator can be used on functions which return objects
       intended to be async context managers. The object returned by
       the function should implement __aenter__ and __aexit__ methods
       to run when the async context is entered and exited.

       This wrapper also allows the use of "await" on the function being
       decorated, to return the context manager without entering it.

    """

    @functools.wraps(coro)
    def context_wrapper(*args, **kwargs) -> _ACMWrapper[_ACM]:
        """Return an async context manager wrapper for this coroutine"""

        return _ACMWrapper(coro(*args, **kwargs))

    return cast(_ACMWrapperFunc[_ACM], context_wrapper)

This avoids the errors above related to __aenter__ and __aexit__ not being present. However, it's not able to catch errors in the arguments passed in to the coroutine which is decorated, since _ACMWrapperFunc takes arbitrary arguments. What I'd like to do here is bind the arguments to exactly match the coro being decorated, but still have the return type be an _ACMWrapper rather than the plain Awaitable that it was originally. It feels like this should be doable with a type variable, but I haven't quite figured out how. Callable doesn't seem to allow a type variable as the first indexed argument. It must be a list of types or .... Any suggestions? I also considered trying to use a Protocol which defines __call__ instead of Callable, but I still don't know how to specify the arguments in that as a type variable.

[JelleZijlstra]

Mypy has enough support for ParamSpec (in recent versions) that it will not crash but just degenerate to Any.

[ronf]

Terrific, thanks! I'll install "pyright" and give it a try...

[ronf]

I think I've managed to get ParamSpec working in pyright, but in mypy 0.910 (which appears to still be the latest version), I'm getting:

asyncssh/misc.py:280: error: The first argument to Callable must be a list of types or "..."
asyncssh/misc.py:281: error: The first argument to Callable must be a list of types or "..."

The lines in question are:

_ACMCoroFunc = Callable[_ACMParams, Awaitable[_ACM]]
_ACMWrapperFunc = Callable[_ACMParams, _ACMWrapper[_ACM]]

where _ACMParams is defined as:

_ACMParams = ParamSpec('_ACMParams')

Any idea how to avoid the mypy error here?

[JelleZijlstra]

# type: ignore is the answer for now, unfortunately.

[ronf]

So, I thought I had ParamSpec working, but upon closer inspection I was wrong. I've been poking at this for a few days now and can't seem to figure out the right set of annotations to make pyright happy. It doesn't seem to be picking up the generic return type properly in one place, reporting it as 'Unknown':

Here's the decorator and some test code for it, which runs fine but gets complaints from pyright:

import asyncio
import functools
from types import TracebackType
from typing import Any, AsyncContextManager, Awaitable, Callable, Generator
from typing import Generic, Optional, ParamSpec, Type, TypeVar

R = TypeVar('R', bound=AsyncContextManager)

class _Wrapper(Generic[R]):
    def __init__(self, coro: Awaitable[R]):
        self._coro = coro
        self._coro_result: Optional[R] = None

    def __await__(self) -> Generator[Any, None, R]:
        return self._coro.__await__()

    async def __aenter__(self) -> R:
        self._coro_result = await self._coro

        return await self._coro_result.__aenter__()

    async def __aexit__(self, exc_type: Type[BaseException],
                        exc_value: BaseException,
                        traceback: TracebackType) -> Optional[bool]:
        assert self._coro_result is not None

        exit_result = await self._coro_result.__aexit__(
            exc_type, exc_value, traceback)

        self._coro_result = None

        return exit_result

P = ParamSpec('P')
AF = Callable[P, Awaitable[R]]
WF = Callable[P, _Wrapper[R]]

def async_context_manager(coro: AF[P, R]) -> WF[P, R]:
    @functools.wraps(coro)
    def wrapfunc(*args: P.args, **kwargs: P.kwargs) -> _Wrapper[R]:
        return _Wrapper(coro(*args, **kwargs))

    return wrapfunc

class MyCtxMgr:
    def __init__(self, x: int):
        self.x = x
        print('Created: x =', x)

    async def __aenter__(self) -> 'MyCtxMgr':
        print('Entered context, x =', self.x)
        return self

    async def __aexit__(self, exc_type: Type[BaseException],
                        exc_value: BaseException,
                        traceback: TracebackType) -> Optional[bool]:
        print('Exited context, x =', self.x)
        return False

@async_context_manager
async def get_context(x: int) -> MyCtxMgr:
    return MyCtxMgr(x)

async def run():
    ctx = await get_context(1)

    async with ctx:
        pass

    async with get_context(2):
        pass

asyncio.run(run())

Here's the error I'm getting:

npx: installed 1 in 1.425s
No configuration file found.
No pyproject.toml file found.
stubPath /private/tmp/typings is not a valid directory.
Assuming Python platform Darwin
Searching for source files
File or directory "/private/tmp/--" does not exist.
Found 1 source file
/private/tmp/t.py
  /private/tmp/t.py:60:2 - error: Argument of type "(x: int) -> Coroutine[Any, Any, MyCtxMgr]" cannot be assigned to parameter "coro" of type "AF[P@deco, R@deco]" in function "deco"
    Type "(x: int) -> Coroutine[Any, Any, MyCtxMgr]" cannot be assigned to type "AF[P@deco, R@deco]"
      Function return type "Coroutine[Any, Any, MyCtxMgr]" is incompatible with type "Awaitable[R@deco]"
        TypeVar "_T_co@Awaitable" is covariant
          Type "AsyncContextManager[Unknown]" cannot be assigned to type "MyCtxMgr" (reportGeneralTypeIssues)
1 error, 0 warnings, 0 infos
Completed in 0.556sec

The line it is flagging is the @async_context_manager line where the decorator is run. I've tried defining the TypeVar R both as shown here and with covariant=True, but neither one works. When setting covariant=True on R, I get back a slightly different error:

npx: installed 1 in 1.61s
No configuration file found.
No pyproject.toml file found.
stubPath /private/tmp/typings is not a valid directory.
Assuming Python platform Darwin
Searching for source files
File or directory "/private/tmp/--" does not exist.
Found 1 source file
/private/tmp/t.py
  /private/tmp/t.py:60:2 - error: Argument of type "(x: int) -> Coroutine[Any, Any, MyCtxMgr]" cannot be assigned to parameter "coro" of type "AF[P@async_context_manager, R@async_context_manager]" in function "async_context_manager"
    Type "(x: int) -> Coroutine[Any, Any, MyCtxMgr]" cannot be assigned to type "AF[P@async_context_manager, R@async_context_manager]"
      Function return type "Coroutine[Any, Any, MyCtxMgr]" is incompatible with type "Awaitable[R@async_context_manager]"
        TypeVar "_T_co@Awaitable" is covariant
          Type "AsyncContextManager[Unknown]" cannot be assigned to type "MyCtxMgr" (reportGeneralTypeIssues)
1 error, 0 warnings, 0 infos
Completed in 0.559sec

The thing that jumps out at me here in both cases is the AsyncContextManager[Unknown] rather than the expected return type, but I don't know how to fix this. Do I need to do something different with the "bound" argument in the TypeVar definition of R? I need the binding here because the coroutine passed into the decorator needs to already have __aenter__ and __aexit__ methods defined on it.

[ronf]

I got a chance to look at this again today, and found the issue! It turns out that the definition of __aexit__ in both _Wrapper and _MyCtxMgr was slightly off, missing "Optional" on some of the arguments, and thus making those classes incompatible with the binding on R to match the protocol AsyncContextManager. Here's a corrected version:

import asyncio
import functools
from types import TracebackType
from typing import Any, AsyncContextManager, Awaitable, Callable, Generator
from typing import Generic, Optional, ParamSpec, Type, TypeVar

R = TypeVar('R', bound=AsyncContextManager)

class _Wrapper(Generic[R]):
    def __init__(self, coro: Awaitable[R]):
        self._coro = coro
        self._coro_result: Optional[R] = None

    def __await__(self) -> Generator[Any, None, R]:
        return self._coro.__await__()

    async def __aenter__(self) -> R:
        self._coro_result = await self._coro

        return await self._coro_result.__aenter__()

    async def __aexit__(self, exc_type: Optional[Type[BaseException]],
                        exc_value: Optional[BaseException],
                        traceback: Optional[TracebackType]) -> Optional[bool]:
        assert self._coro_result is not None

        exit_result = await self._coro_result.__aexit__(
            exc_type, exc_value, traceback)

        self._coro_result = None

        return exit_result

P = ParamSpec('P')
AF = Callable[P, Awaitable[R]]
WF = Callable[P, _Wrapper[R]]

def async_context_manager(coro: AF[P, R]) -> WF[P, R]:
    @functools.wraps(coro)
    def wrapfunc(*args: P.args, **kwargs: P.kwargs) -> _Wrapper[R]:
        return _Wrapper(coro(*args, **kwargs))

    return wrapfunc

class MyCtxMgr:
    def __init__(self, x: int):
        self.x = x
        print('Created: x =', x)

    async def __aenter__(self) -> 'MyCtxMgr':
        print('Entered context, x =', self.x)
        return self

    async def __aexit__(self, exc_type: Optional[Type[BaseException]],
                        exc_value: Optional[BaseException],
                        traceback: Optional[TracebackType]) -> Optional[bool]:
        print('Exited context, x =', self.x)
        return False

@async_context_manager
async def get_context(x: int) -> MyCtxMgr:
    return MyCtxMgr(x)

async def run():
    ctx = await get_context(1)

    async with ctx:
        pass

    async with get_context(2):
        pass

asyncio.run(run())

With this version, if I change the arguments to something incorrect, pyright is able to catch it. For instance:

    ctx = await get_context(1, 'xxx')

results in:

  /private/tmp/t1.py:65:32 - error: Expected 1 positional argument (reportGeneralTypeIssues)
1 error, 0 warnings, 0 infos

Similarly:

    ctx = await get_context(1, key='xxx')

results in:

  /private/tmp/t1.py:65:32 - error: No parameter named "key" (reportGeneralTypeIssues)
1 error, 0 warnings, 0 infos

Unfortunately, mypy doesn't like the type aliases of AF and WF using P as an argument. To make mypy happy, these definitions need to be removed and expanded to something like:

def async_context_manager(coro: Callable[P, Awaitable[R]]) -> \
        Callable[P, _Wrapper[R]]:

This use of P seems to be accepted by mypy (but basically ignored). It also reports errors later when referencing P.args and P.kwargs, so the following additional change is necessary to allow the code to work with both mypy and pyright:

    def wrapfunc(*args: P.args,                      # type: ignore
                 **kwargs: P.kwargs) -> _Wrapper[R]: # type: ignore