More tweaks for type completion

Author: mikeshardmind

src/async_utils/_internal_types.py

@@ -30,9 +30,10 @@
 if TYPE_CHECKING:
     from typing import Any, Concatenate, Literal, Never, Protocol, Self
 else:
+    Protocol = object
 
     def __getattr__(name: str):
-        if name in {"Any", "Concatenate", "Literal", "Never", "Protocol", "Self"}:
+        if name in {"Any", "Concatenate", "Literal", "Never", "Self"}:
             import typing
 
             return getattr(typing, name)

src/async_utils/_typings.py

@@ -16,16 +16,29 @@
 
 import asyncio
 import concurrent.futures as cf
-from collections.abc import Hashable
+from collections.abc import Awaitable, Callable, Coroutine, Hashable
 from functools import partial, wraps
 
-from . import _internal_types as i_types
 from . import _typings as t
 from ._paramkey import make_key
 from .lru import LRU
 
 __all__ = ("corocache", "lrucorocache")
 
+type CoroFunc[**P, R] = Callable[P, Coroutine[t.Any, t.Any, R]]
+type CoroLike[**P, R] = Callable[P, Awaitable[R]]
+
+#: Note CacheTransformers recieve a tuple (args) and dict(kwargs)
+#: rather than a ParamSpec of the decorated function.
+#: Warning: Mutations will impact callsite, return new objects as needed.
+type CacheTransformer = Callable[
+    [tuple[t.Any, ...], dict[str, t.Any]], tuple[tuple[t.Any, ...], dict[str, t.Any]]
+]
+
+
+class CoroCacheDeco(t.Protocol):
+    def __call__[**P, R](self, c: CoroLike[P, R], /) -> CoroFunc[P, R]: ...
+
 
 def _chain_fut[R](c_fut: cf.Future[R], a_fut: asyncio.Future[R]) -> None:
     if a_fut.cancelled():
@@ -39,8 +52,8 @@ def _chain_fut[R](c_fut: cf.Future[R], a_fut: asyncio.Future[R]) -> None:
 def corocache(
     ttl: float | None = None,
     *,
-    cache_transform: i_types.CacheTransformer | None = None,
-) -> i_types.CoroCacheDeco:
+    cache_transform: CacheTransformer | None = None,
+) -> CoroCacheDeco:
     """Cache the results of the decorated coroutine.
 
     This is less powerful than the version in task_cache.py but may work better
@@ -73,7 +86,7 @@ def corocache(
         def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any], /) -> Hashable:
             return make_key(*cache_transform(args, kwds))
 
-    def wrapper[**P, R](coro: i_types.CoroLike[P, R], /) -> i_types.CoroFunc[P, R]:
+    def wrapper[**P, R](coro: CoroLike[P, R], /) -> CoroFunc[P, R]:
         internal_cache: dict[Hashable, cf.Future[R]] = {}
         internal_taskset: set[asyncio.Task[R]] = set()
 
@@ -112,8 +125,8 @@ def lrucorocache(
     ttl: float | None = None,
     maxsize: int = 1024,
     *,
-    cache_transform: i_types.CacheTransformer | None = None,
-) -> i_types.CoroCacheDeco:
+    cache_transform: CacheTransformer | None = None,
+) -> CoroCacheDeco:
     """Cache the results of the decorated coroutine.
 
     This is less powerful than the version in task_cache.py but may work better
@@ -152,7 +165,7 @@ def lrucorocache(
         def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any], /) -> Hashable:
             return make_key(*cache_transform(args, kwds))
 
-    def wrapper[**P, R](coro: i_types.CoroLike[P, R], /) -> i_types.CoroFunc[P, R]:
+    def wrapper[**P, R](coro: CoroLike[P, R], /) -> CoroFunc[P, R]:
         internal_cache: LRU[Hashable, cf.Future[R]] = LRU(maxsize)
         internal_taskset: set[asyncio.Task[R]] = set()
 

src/async_utils/corofunc_cache.py

@@ -16,16 +16,30 @@
 
 import asyncio
 import concurrent.futures as cf
-from collections.abc import Callable, Hashable
+from collections.abc import Callable, Coroutine, Hashable
 from functools import partial, wraps
 
-from . import _internal_types as i_types
 from . import _typings as t
 from ._paramkey import make_key
 from .lru import LRU
 
 __all__ = ("lrutaskcache", "taskcache")
 
+type CoroFunc[**P, R] = Callable[P, Coroutine[t.Any, t.Any, R]]
+type TaskFunc[**P, R] = Callable[P, asyncio.Task[R]]
+type TaskCoroFunc[**P, R] = CoroFunc[P, R] | TaskFunc[P, R]
+
+#: Note CacheTransformers recieve a tuple (args) and dict(kwargs)
+#: rather than a ParamSpec of the decorated function.
+#: Warning: Mutations will impact callsite, return new objects as needed.
+type CacheTransformer = Callable[
+    [tuple[t.Any, ...], dict[str, t.Any]], tuple[tuple[t.Any, ...], dict[str, t.Any]]
+]
+
+
+class TaskCacheDeco(t.Protocol):
+    def __call__[**P, R](self, c: TaskCoroFunc[P, R], /) -> TaskFunc[P, R]: ...
+
 
 # Non-annotation assignments for transformed functions
 _WRAP_ASSIGN = ("__module__", "__name__", "__qualname__", "__doc__")
@@ -44,7 +58,7 @@ class _WrappedSignature[**P, R]:
     #: PYUPGRADE: Ensure inspect.signature still accepts this
     # as func.__signature__
     # Known working: py 3.12.0 - py3.14a6 range inclusive
-    def __init__(self, f: i_types.TaskCoroFunc[P, R], w: i_types.TaskFunc[P, R]) -> None:
+    def __init__(self, f: TaskCoroFunc[P, R], w: TaskFunc[P, R]) -> None:
         self._f: Callable[..., t.Any] = f  # anotation needed for inspect use below....
         self._w = w
         self._sig: t.Any | None = None
@@ -92,8 +106,8 @@ async def _await[R](fut: asyncio.Future[R]) -> R:
 def taskcache(
     ttl: float | None = None,
     *,
-    cache_transform: i_types.CacheTransformer | None = None,
-) -> i_types.TaskCacheDeco:
+    cache_transform: CacheTransformer | None = None,
+) -> TaskCacheDeco:
     """Cache the results of the decorated coroutine.
 
     Decorator to modify coroutine functions to instead act as functions
@@ -132,7 +146,7 @@ def taskcache(
         def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any], /) -> Hashable:
             return make_key(*cache_transform(args, kwds))
 
-    def wrapper[**P, R](coro: i_types.TaskCoroFunc[P, R], /) -> i_types.TaskFunc[P, R]:
+    def wrapper[**P, R](coro: TaskCoroFunc[P, R], /) -> TaskFunc[P, R]:
         internal_cache: dict[Hashable, cf.Future[R]] = {}
 
         def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:
@@ -171,8 +185,8 @@ def lrutaskcache(
     ttl: float | None = None,
     maxsize: int = 1024,
     *,
-    cache_transform: i_types.CacheTransformer | None = None,
-) -> i_types.TaskCacheDeco:
+    cache_transform: CacheTransformer | None = None,
+) -> TaskCacheDeco:
     """Cache the results of the decorated coroutine.
 
     Decorator to modify coroutine functions to instead act as functions
@@ -218,7 +232,7 @@ def lrutaskcache(
         def key_func(args: tuple[t.Any, ...], kwds: dict[t.Any, t.Any], /) -> Hashable:
             return make_key(*cache_transform(args, kwds))
 
-    def wrapper[**P, R](coro: i_types.TaskCoroFunc[P, R], /) -> i_types.TaskFunc[P, R]:
+    def wrapper[**P, R](coro: TaskCoroFunc[P, R], /) -> TaskFunc[P, R]:
         internal_cache: LRU[Hashable, cf.Future[R]] = LRU(maxsize)
 
         def _internal_cache_evict(key: Hashable, _ignored_task: object) -> None:

src/async_utils/task_cache.py

@@ -23,7 +23,7 @@
 
 from . import _typings as t
 
-log = logging.getLogger(__name__)
+_log = logging.getLogger(__name__)
 
 __all__ = ("Waterfall",)
 
@@ -117,9 +117,9 @@ def put(self, item: T) -> None:
 
     def _user_done_callback(self, num: int, future: asyncio.Future[t.Any]) -> None:
         if future.cancelled():
-            log.warning("Callback cancelled due to timeout")
+            _log.warning("Callback cancelled due to timeout")
         elif exc := future.exception():
-            log.error("Exception in user callback", exc_info=exc)
+            _log.error("Exception in user callback", exc_info=exc)
 
         for _ in range(num):
             self.queue.task_done()