Fix .pyi typing for asyncio (#134)

asynq/decorators.py

@@ -15,7 +15,7 @@
 
 import asyncio
 import inspect
-from typing import Any, Awaitable
+from typing import Any, Coroutine
 
 import qcore.decorators
 import qcore.helpers as core_helpers
@@ -133,7 +133,7 @@ def _fn_wrapper(self, args, kwargs):
         return
         yield
 
-    def asyncio(self, *args, **kwargs) -> Awaitable[Any]:
+    def asyncio(self, *args, **kwargs) -> Coroutine[Any, Any, Any]:
         if self.asyncio_fn is None:
             if inspect.isgeneratorfunction(self.fn):
 
@@ -184,7 +184,7 @@ def asynq(self, *args, **kwargs):
         else:
             return self.decorator.asynq(self.instance, *args, **kwargs)
 
-    def asyncio(self, *args, **kwargs) -> Awaitable[Any]:
+    def asyncio(self, *args, **kwargs) -> Coroutine[Any, Any, Any]:
         if self.instance is None:
             return self.decorator.asyncio(*args, **kwargs)
         else:

asynq/decorators.pyi

@@ -1,7 +1,7 @@
 from typing import (
     Any,
-    Awaitable,
     Callable,
+    Coroutine,
     Generic,
     Mapping,
     Optional,
@@ -16,6 +16,8 @@ import qcore.decorators
 from . import async_task, futures
 
 _T = TypeVar("_T")
+_Coroutine = Coroutine[Any, Any, Any]
+_CoroutineFn = Callable[..., _Coroutine]
 
 def lazy(fn: Callable[..., _T]) -> Callable[..., futures.FutureBase[_T]]: ...
 def has_async_fn(fn: object) -> bool: ...
@@ -36,16 +38,19 @@ class PureAsyncDecorator(qcore.decorators.DecoratorBase, Generic[_T]):
         fn: Callable[..., _T],
         task_cls: Optional[Type[futures.FutureBase]],
         kwargs: Mapping[str, Any] = ...,
+        asyncio_fn: Optional[_CoroutineFn] = ...,
     ) -> None: ...
     def name(self) -> str: ...
     def is_pure_async_fn(self) -> bool: ...
+    def asyncio(self, *args, **kwargs) -> _Coroutine: ...
     def __call__(
         self, *args: Any, **kwargs: Any
     ) -> Union[_T, futures.FutureBase[_T]]: ...
     def __get__(self, owner: Any, cls: Any) -> PureAsyncDecorator[_T]: ...  # type: ignore
 
 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 AsyncDecorator(PureAsyncDecorator[_T]):
     binder_cls = AsyncDecoratorBinder  # type: ignore
@@ -54,11 +59,10 @@ class AsyncDecorator(PureAsyncDecorator[_T]):
         fn: Callable[..., _T],
         cls: Optional[Type[futures.FutureBase]],
         kwargs: Mapping[str, Any] = ...,
-        asyncio_fn: Callable[..., Awaitable] = ...,
+        asyncio_fn: Optional[_CoroutineFn] = ...,
     ): ...
     def is_pure_async_fn(self) -> bool: ...
     def asynq(self, *args: Any, **kwargs: Any) -> async_task.AsyncTask[_T]: ...
-    def asyncio(self, *args, **kwargs) -> Awaitable[Any]: ...
     def __call__(self, *args: Any, **kwargs: Any) -> _T: ...
     def __get__(self, owner: Any, cls: Any) -> AsyncDecorator[_T]: ...  # type: ignore
 
@@ -77,11 +81,18 @@ class AsyncAndSyncPairDecorator(AsyncDecorator[_T]):
     def __get__(self, owner: Any, cls: Any) -> Any: ...
 
 class AsyncProxyDecorator(AsyncDecorator[_T]):
-    def __init__(self, fn: Callable[..., futures.FutureBase[_T]]) -> None: ...
+    def __init__(
+        self,
+        fn: Callable[..., futures.FutureBase[_T]],
+        asyncio_fn: Optional[_CoroutineFn] = ...,
+    ) -> None: ...
 
 class AsyncAndSyncPairProxyDecorator(AsyncProxyDecorator[_T]):
     def __init__(
-        self, fn: Callable[..., futures.FutureBase[_T]], sync_fn: Callable[..., _T]
+        self,
+        fn: Callable[..., futures.FutureBase[_T]],
+        sync_fn: Callable[..., _T],
+        asyncio_fn: Optional[_CoroutineFn] = ...,
     ) -> None: ...
     def __call__(self, *args: Any, **kwargs: Any) -> _T: ...
 
@@ -97,23 +108,28 @@ def asynq(  # type: ignore
     *,
     sync_fn: Optional[Callable[..., Any]] = ...,
     cls: Type[futures.FutureBase] = ...,
+    asyncio_fn: Optional[_CoroutineFn] = ...,
     **kwargs: Any,
 ) -> _MkAsyncDecorator: ...
 @overload
 def asynq(
     pure: bool,
     sync_fn: Optional[Callable[..., Any]] = ...,
     cls: Type[futures.FutureBase] = ...,
-    asyncio_fn: Callable[..., Awaitable] = ...,
+    asyncio_fn: Optional[_CoroutineFn] = ...,
     **kwargs: Any,
 ) -> _MkPureAsyncDecorator: ...  # type: ignore
 @overload
 def async_proxy(
-    *, sync_fn: Optional[Callable[..., Any]] = ...
+    *,
+    sync_fn: Optional[Callable[..., Any]] = ...,
+    asyncio_fn: Optional[_CoroutineFn] = ...,
 ) -> _MkAsyncDecorator: ...
 @overload
 def async_proxy(
-    pure: bool, sync_fn: Optional[Callable[..., Any]] = ...
+    pure: bool,
+    sync_fn: Optional[Callable[..., Any]] = ...,
+    asyncio_fn: Optional[_CoroutineFn] = ...,
 ) -> _MkPureAsyncDecorator: ...
 @asynq()
 def async_call(