Be more permissive

Author: mikeshardmind

async_utils/__init__.py

@@ -13,4 +13,4 @@
 #   limitations under the License.
 
 
-__version__ = "7.0.3"
+__version__ = "7.0.4"

async_utils/corofunc_cache.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable, Coroutine, Hashable
+from collections.abc import Awaitable, Callable, Coroutine, Hashable
 from functools import partial, wraps
 from typing import Any, ParamSpec, TypeVar
 
@@ -26,11 +26,10 @@
 
 P = ParamSpec("P")
 T = TypeVar("T")
-K = TypeVar("K")
-V = TypeVar("V")
 
 
 type CoroFunc[**P, T] = Callable[P, Coroutine[Any, Any, T]]
+type CoroLike[**P, T] = Callable[P, Awaitable[T]]
 
 
 class LRU[K, V]:
@@ -59,7 +58,7 @@ def remove(self, key: K) -> None:
 
 def corocache(
     ttl: float | None = None,
-) -> Callable[[CoroFunc[P, T]], CoroFunc[P, T]]:
+) -> Callable[[CoroLike[P, T]], CoroFunc[P, T]]:
     """Decorator to cache coroutine functions.
 
     This is less powerful than the version in task_cache.py but may work better for
@@ -71,37 +70,37 @@ def corocache(
 
     The ordering of args and kwargs matters."""
 
-    def wrapper(coro: CoroFunc[P, T]) -> CoroFunc[P, T]:
-        internal_cache: dict[Hashable, asyncio.Task[T]] = {}
+    def wrapper(coro: CoroLike[P, T]) -> CoroFunc[P, T]:
+        internal_cache: dict[Hashable, asyncio.Future[T]] = {}
 
         async def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
             key = make_key(args, kwargs)
             try:
                 return await internal_cache[key]
             except KeyError:
-                internal_cache[key] = task = asyncio.create_task(coro(*args, **kwargs))
+                internal_cache[key] = fut = asyncio.ensure_future(coro(*args, **kwargs))
                 if ttl is not None:
-                    # This results in internal_cache.pop(key, task) later
+                    # This results in internal_cache.pop(key, fut) later
                     # while avoiding a late binding issue with a lambda instead
                     call_after_ttl = partial(
                         asyncio.get_running_loop().call_later,
                         ttl,
                         internal_cache.pop,
                         key,
                     )
-                    task.add_done_callback(call_after_ttl)
-                return await task
+                    fut.add_done_callback(call_after_ttl)
+                return await fut
 
         return wrapped
 
     return wrapper
 
 
-def _lru_evict(ttl: float, cache: LRU[Hashable, Any], key: Hashable, _ignored_task: object) -> None:
+def _lru_evict(ttl: float, cache: LRU[Hashable, Any], key: Hashable, _ignored_fut: object) -> None:
     asyncio.get_running_loop().call_later(ttl, cache.remove, key)
 
 
-def lrucorocache(ttl: float | None = None, maxsize: int = 1024) -> Callable[[CoroFunc[P, T]], CoroFunc[P, T]]:
+def lrucorocache(ttl: float | None = None, maxsize: int = 1024) -> Callable[[CoroLike[P, T]], CoroFunc[P, T]]:
     """Decorator to cache coroutine functions.
 
     This is less powerful than the version in task_cache.py but may work better for
@@ -113,22 +112,22 @@ def lrucorocache(ttl: float | None = None, maxsize: int = 1024) -> Callable[[Cor
 
     The ordering of args and kwargs matters.
 
-    tasks are evicted by LRU and ttl.
+    futs are evicted by LRU and ttl.
     """
 
-    def wrapper(coro: CoroFunc[P, T]) -> CoroFunc[P, T]:
-        internal_cache: LRU[Hashable, asyncio.Task[T]] = LRU(maxsize)
+    def wrapper(coro: CoroLike[P, T]) -> CoroFunc[P, T]:
+        internal_cache: LRU[Hashable, asyncio.Future[T]] = LRU(maxsize)
 
         @wraps(coro)
         async def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
             key = make_key(args, kwargs)
             try:
                 return await internal_cache[key]
             except KeyError:
-                internal_cache[key] = task = asyncio.create_task(coro(*args, **kwargs))
+                internal_cache[key] = fut = asyncio.ensure_future(coro(*args, **kwargs))
                 if ttl is not None:
-                    task.add_done_callback(partial(_lru_evict, ttl, internal_cache, key))
-                return await task
+                    fut.add_done_callback(partial(_lru_evict, ttl, internal_cache, key))
+                return await fut
 
         return wrapped
 

async_utils/task_cache.py

@@ -27,6 +27,11 @@
 P = ParamSpec("P")
 T = TypeVar("T")
 
+# Use below doesn't accept non-task Futures, so can't accept general awaitables
+type CoroFunc[**P, T] = Callable[P, Coroutine[Any, Any, T]]
+type TaskFunc[**P, T] = CoroFunc[P, T] | Callable[P, asyncio.Task[T]]
+type TaskCoroFunc[**P, T] = CoroFunc[P, T] | TaskFunc[P, T]
+
 
 class LRU[K, V]:
     def __init__(self, maxsize: int, /):
@@ -54,7 +59,7 @@ def remove(self, key: K) -> None:
 
 def taskcache(
     ttl: float | None = None,
-) -> Callable[[Callable[P, Coroutine[Any, Any, T]]], Callable[P, asyncio.Task[T]]]:
+) -> Callable[[TaskCoroFunc[P, T]], TaskFunc[P, T]]:
     """Decorator to modify coroutine functions to instead act as functions returning cached tasks.
 
     For general use, this leaves the end user API largely the same,
@@ -66,7 +71,7 @@ def taskcache(
 
     The ordering of args and kwargs matters."""
 
-    def wrapper(coro: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, asyncio.Task[T]]:
+    def wrapper(coro: TaskCoroFunc[P, T]) -> TaskFunc[P, T]:
         internal_cache: dict[Hashable, asyncio.Task[T]] = {}
 
         @wraps(coro, assigned=("__module__", "__name__", "__qualname__", "__doc__"))
@@ -75,7 +80,7 @@ def wrapped(*args: P.args, **kwargs: P.kwargs) -> asyncio.Task[T]:
             try:
                 return internal_cache[key]
             except KeyError:
-                internal_cache[key] = task = asyncio.create_task(coro(*args, **kwargs))
+                internal_cache[key] = task = asyncio.ensure_future(coro(*args, **kwargs))
                 if ttl is not None:
                     # This results in internal_cache.pop(key, task) later
                     # while avoiding a late binding issue with a lambda instead
@@ -97,9 +102,7 @@ def _lru_evict(ttl: float, cache: LRU[Hashable, Any], key: Hashable, _ignored_ta
     asyncio.get_running_loop().call_later(ttl, cache.remove, key)
 
 
-def lrutaskcache(
-    ttl: float | None = None, maxsize: int = 1024
-) -> Callable[[Callable[P, Coroutine[Any, Any, T]]], Callable[P, asyncio.Task[T]]]:
+def lrutaskcache(ttl: float | None = None, maxsize: int = 1024) -> Callable[[TaskCoroFunc[P, T]], TaskFunc[P, T]]:
     """Decorator to modify coroutine functions to instead act as functions returning cached tasks.
 
     For general use, this leaves the end user API largely the same,
@@ -114,7 +117,7 @@ def lrutaskcache(
     tasks are evicted by LRU and ttl.
     """
 
-    def wrapper(coro: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, asyncio.Task[T]]:
+    def wrapper(coro: TaskCoroFunc[P, T]) -> TaskFunc[P, T]:
         internal_cache: LRU[Hashable, asyncio.Task[T]] = LRU(maxsize)
 
         @wraps(coro, assigned=("__module__", "__name__", "__qualname__", "__doc__"))
@@ -123,7 +126,7 @@ def wrapped(*args: P.args, **kwargs: P.kwargs) -> asyncio.Task[T]:
             try:
                 return internal_cache[key]
             except KeyError:
-                internal_cache[key] = task = asyncio.create_task(coro(*args, **kwargs))
+                internal_cache[key] = task = asyncio.ensure_future(coro(*args, **kwargs))
                 if ttl is not None:
                     task.add_done_callback(partial(_lru_evict, ttl, internal_cache, key))
                 return task