From 4cdc7a669cd5789944c149ca902b56262af5ee14 Mon Sep 17 00:00:00 2001 From: dkang-quora <135926539+dkang-quora@users.noreply.github.com> Date: Wed, 17 Jul 2024 07:06:55 +0900 Subject: [PATCH] Further improve typing (#145) Co-authored-by: Jelle Zijlstra --- asynq/decorators.pyi | 82 ++++++++++++++++++++++++++------------ asynq/tests/test_typing.py | 31 +++++++++++--- asynq/tools.pyi | 12 ++++-- 3 files changed, 90 insertions(+), 35 deletions(-) diff --git a/asynq/decorators.pyi b/asynq/decorators.pyi index da579c5..bbf7dfb 100644 --- a/asynq/decorators.pyi +++ b/asynq/decorators.pyi @@ -1,8 +1,8 @@ from typing import ( Any, - Awaitable, Callable, Coroutine, + Generator, Generic, Mapping, Optional, @@ -12,14 +12,16 @@ from typing import ( ) import qcore.decorators -from typing_extensions import ParamSpec +from typing_extensions import Concatenate, Literal, ParamSpec from . import async_task, futures _P = ParamSpec("_P") +_P2 = ParamSpec("_P2") _T = TypeVar("_T") -_Coroutine = Coroutine[Any, Any, Any] -_CoroutineFn = Callable[..., _Coroutine] +_T2 = TypeVar("_T2") +_G = Generator[Any, Any, _T] # Generator that returns _T +_Coroutine = Coroutine[Any, Any, _T] def lazy(fn: Callable[_P, _T]) -> Callable[_P, futures.FutureBase[_T]]: ... def has_async_fn(fn: object) -> bool: ... @@ -30,17 +32,28 @@ def get_async_fn( ) -> Optional[Callable[..., futures.FutureBase[Any]]]: ... def get_async_or_sync_fn(fn: object) -> Any: ... -class PureAsyncDecoratorBinder(qcore.decorators.DecoratorBinder): +class PureAsyncDecoratorBinder(qcore.decorators.DecoratorBinder[_T], Generic[_T, _P]): def is_pure_async_fn(self) -> bool: ... + def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T: ... class PureAsyncDecorator(qcore.decorators.DecoratorBase, Generic[_T, _P]): binder_cls = PureAsyncDecoratorBinder + + @overload def __init__( self, fn: Callable[_P, Any], # TODO overloads for Generator[Any, Any, _T] and _T task_cls: Optional[type[futures.FutureBase]], kwargs: Mapping[str, Any] = ..., - asyncio_fn: Optional[Callable[_P, Awaitable[_T]]] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., + ) -> None: ... + @overload + def __init__( + self, + fn: Callable[_P, Generator[Any, Any, _T]], + task_cls: Optional[type[futures.FutureBase]], + kwargs: Mapping[str, Any] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., ) -> None: ... def name(self) -> str: ... def is_pure_async_fn(self) -> bool: ... @@ -50,27 +63,42 @@ class PureAsyncDecorator(qcore.decorators.DecoratorBase, Generic[_T, _P]): def __call__( self, *args: Any, **kwargs: Any ) -> Union[_T, futures.FutureBase[_T]]: ... - def __get__(self, owner: Any, cls: Any) -> PureAsyncDecorator[_T, _P]: ... # type: ignore[override] + def __get__( + self: PureAsyncDecorator[_T2, Concatenate[Any, _P2]], owner: Any, cls: Any + ) -> PureAsyncDecoratorBinder[_T2, _P2]: ... -class AsyncDecoratorBinder(qcore.decorators.DecoratorBinder, Generic[_T]): - def asynq(self, *args: Any, **kwargs: Any) -> async_task.AsyncTask[_T]: ... - def asyncio(self, *args, **kwargs) -> _Coroutine: ... +class AsyncDecoratorBinder(qcore.decorators.DecoratorBinder, Generic[_T, _P]): + def asynq( + self, *args: _P.args, **kwargs: _P.kwargs + ) -> async_task.AsyncTask[_T]: ... + def asyncio( + self, *args: _P.args, **kwargs: _P.kwargs + ) -> Coroutine[Any, Any, _T]: ... class AsyncDecorator(PureAsyncDecorator[_T, _P]): binder_cls = AsyncDecoratorBinder # type: ignore + @overload def __init__( self, fn: Callable[_P, Any], # TODO overloads for Generator[Any, Any, _T] and _T cls: Optional[type[futures.FutureBase]], kwargs: Mapping[str, Any] = ..., - asyncio_fn: Optional[Callable[_P, Awaitable[_T]]] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., + ): ... + @overload + def __init__( + self, + fn: Callable[_P, Generator[Any, Any, _T]], + cls: Optional[type[futures.FutureBase]], + kwargs: Mapping[str, Any] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., ): ... def is_pure_async_fn(self) -> bool: ... def asynq(self, *args: Any, **kwargs: Any) -> async_task.AsyncTask[_T]: ... def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _T: ... - def __get__(self, owner: Any, cls: Any) -> AsyncDecorator[_T, _P]: ... # type: ignore[override] + def __get__(self: PureAsyncDecorator[_T2, Concatenate[Any, _P2]], owner: Any, cls: Any) -> AsyncDecoratorBinder[_T2, _P2]: ... # type: ignore[override] -class AsyncAndSyncPairDecoratorBinder(AsyncDecoratorBinder[_T]): ... +class AsyncAndSyncPairDecoratorBinder(AsyncDecoratorBinder[_T, _P]): ... class AsyncAndSyncPairDecorator(AsyncDecorator[_T, _P]): binder_cls = AsyncAndSyncPairDecoratorBinder # type: ignore @@ -88,7 +116,7 @@ class AsyncProxyDecorator(AsyncDecorator[_T, _P]): def __init__( self, fn: Callable[..., futures.FutureBase[_T]], - asyncio_fn: Optional[_CoroutineFn] = ..., + asyncio_fn: Optional[Callable[..., Coroutine[Any, Any, _T]]] = ..., ) -> None: ... class AsyncAndSyncPairProxyDecorator(AsyncProxyDecorator[_T, _P]): @@ -96,7 +124,7 @@ class AsyncAndSyncPairProxyDecorator(AsyncProxyDecorator[_T, _P]): self, fn: Callable[..., futures.FutureBase[_T]], sync_fn: Callable[..., _T], - asyncio_fn: Optional[_CoroutineFn] = ..., + asyncio_fn: Optional[Callable[..., Coroutine[Any, Any, _T]]] = ..., ) -> None: ... def __call__(self, *args: Any, **kwargs: Any) -> _T: ... @@ -104,36 +132,38 @@ class _MkAsyncDecorator: def __call__(self, fn: Callable[_P, Any]) -> AsyncDecorator[Any, _P]: ... class _MkPureAsyncDecorator: - def __call__(self, fn: Callable[_P, _T]) -> PureAsyncDecorator[_T, _P]: ... + def __call__(self, fn: Callable[_P, Any]) -> PureAsyncDecorator[Any, _P]: ... # In reality these two can return other Decorator subclasses, but that doesn't matter for callers. @overload -def asynq( # type: ignore +def asynq( *, - sync_fn: Optional[Callable[..., Any]] = ..., + pure: Literal[False] = False, + sync_fn: Optional[Callable[_P, _T]] = ..., cls: type[futures.FutureBase] = ..., - asyncio_fn: Optional[_CoroutineFn] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., **kwargs: Any, ) -> _MkAsyncDecorator: ... @overload def asynq( - pure: bool, - sync_fn: Optional[Callable[..., Any]] = ..., + *, + pure: Literal[True], + sync_fn: Optional[Callable[_P, _T]] = ..., cls: type[futures.FutureBase] = ..., - asyncio_fn: Optional[_CoroutineFn] = ..., + asyncio_fn: Optional[Callable[_P, Coroutine[Any, Any, _T]]] = ..., **kwargs: Any, ) -> _MkPureAsyncDecorator: ... @overload def async_proxy( *, - sync_fn: Optional[Callable[..., Any]] = ..., - asyncio_fn: Optional[_CoroutineFn] = ..., + sync_fn: Optional[Callable[_P, Union[_T, Generator[Any, Any, _T]]]] = ..., + asyncio_fn: Optional[Callable[..., Coroutine[Any, Any, _T]]] = ..., ) -> _MkAsyncDecorator: ... @overload def async_proxy( pure: bool, - sync_fn: Optional[Callable[..., Any]] = ..., - asyncio_fn: Optional[_CoroutineFn] = ..., + sync_fn: Optional[Callable[_P, Union[_T, Generator[Any, Any, _T]]]] = ..., + asyncio_fn: Optional[Callable[..., Coroutine[Any, Any, _T]]] = ..., ) -> _MkPureAsyncDecorator: ... @asynq() def async_call( diff --git a/asynq/tests/test_typing.py b/asynq/tests/test_typing.py index 33fa579..76b1e4b 100644 --- a/asynq/tests/test_typing.py +++ b/asynq/tests/test_typing.py @@ -1,7 +1,10 @@ -from typing import Any, Generator, TYPE_CHECKING +from typing import Any, Generator, TYPE_CHECKING, TypeVar from typing_extensions import assert_type from asynq.decorators import async_call, lazy, asynq from asynq.futures import FutureBase +from asynq.tools import amax + +_T = TypeVar("_T") def test_lazy() -> None: @@ -23,13 +26,25 @@ def generator(x: int) -> Generator[Any, Any, str]: yield None return str(x) - async def caller() -> None: - # This doesn't work, apparently due to a mypy bug - assert_type(await generator.asyncio(1), str) # type: ignore[assert-type] - assert_type(await non_generator.asyncio(1), str) # type: ignore[assert-type] + @asynq() + def generic(x: _T) -> _T: + return x + + class Obj: + @asynq() + def method(self, x: str) -> int: + return int(x) + + async def caller(x: str, obj: Obj) -> None: + assert_type(await generator.asyncio(1), Any) # TODO: str + assert_type(await non_generator.asyncio(1), Any) # TODO: str + assert_type(await generic.asyncio(x), Any) # TODO: str + assert_type(await obj.method.asyncio(x), Any) # TODO: int await non_generator.asyncio() # type: ignore[call-arg] await generator.asyncio() # type: ignore[call-arg] + await generic.asyncio(1, 2) # type: ignore[call-arg] + await obj.method.asyncio(1) # type: ignore[arg-type] def test_async_call() -> None: @@ -39,3 +54,9 @@ def f(x: int) -> str: async_call(f, 1) if TYPE_CHECKING: async_call(f, 1, task_cls=FutureBase) # TODO: this should be an error + + +def test_amax(x: int = 1, y: int = 2) -> None: + if TYPE_CHECKING: + assert_type(amax(1, 2, key=len), Any) # TODO: int + assert_type(amax([1, 2], key=len), Any) # TODO: int diff --git a/asynq/tools.pyi b/asynq/tools.pyi index 6a4d285..40b9fcf 100644 --- a/asynq/tools.pyi +++ b/asynq/tools.pyi @@ -38,16 +38,20 @@ def asorted( ) -> List[_T]: ... @overload @asynq() -def amax(__arg: Iterable[_T], key: Optional[Callable[[_T], Any]] = ...) -> _T: ... +def amax( + arg1: _T, arg2: _T, /, *args: _T, key: Optional[Callable[[_T], Any]] = ... +) -> _T: ... @overload @asynq() -def amax(*args: _T, key: Optional[Callable[[_T], Any]] = ...) -> _T: ... +def amax(arg: Iterable[_T], /, *, key: Optional[Callable[[_T], Any]] = ...) -> _T: ... @overload @asynq() -def amin(__arg: Iterable[_T], key: Optional[Callable[[_T], Any]] = ...) -> _T: ... +def amin( + arg1: _T, arg2: _T, /, *args: _T, key: Optional[Callable[[_T], Any]] = ... +) -> _T: ... @overload @asynq() -def amin(*args: _T, key: Optional[Callable[[_T], Any]] = ...) -> _T: ... +def amin(__arg: Iterable[_T], key: Optional[Callable[[_T], Any]] = ...) -> _T: ... @asynq() def asift( pred: Callable[[_T], bool], items: Iterable[_T]