omnilib/aiosqlite

Better type annotations for contextmanager

andreymal opened this issue · 0 comments

I noticed that contextmanager doesn't preserve signatures of decorated methods:

reveal_type(db.executemany)
# mypy: Revealed type is "def (*Any, **Any) -> aiosqlite.context.Result[aiosqlite.cursor.Cursor]"

Consider using ParamSpec to preserve them:

from typing_extensions import ParamSpec

_P = ParamSpec("_P")

def contextmanager(
    method: Callable[_P, Coroutine[Any, Any, _T]]
) -> Callable[_P, Result[_T]]:
    @wraps(method)
    def wrapper(self, *args: _P.args, **kwargs: _P.kwargs) -> Result[_T]:
        return Result(method(self, *args, **kwargs))

    return wrapper
reveal_type(db.executemany)
# mypy: Revealed type is "def (sql: builtins.str, parameters: typing.Iterable[typing.Iterable[Any]])
#  -> aiosqlite.context.Result[aiosqlite.cursor.Cursor]"