Merge pull request #3 from simonsobs/type

Clean code

pyproject.toml

@@ -79,5 +79,5 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 profile = "black"
 
 [[tool.mypy.overrides]]
-module = ["pluggy.*"]
+module = ["pluggy.*", "apluggy.*"]
 ignore_missing_imports = true

src/apluggy/_wrap/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["PluginManager"]
+
+from .manager import PluginManager

src/apluggy/_wrap/ahook.py

@@ -0,0 +1,19 @@
+import asyncio
+from collections.abc import Callable
+from typing import Any, Coroutine
+
+from pluggy import HookCaller
+from pluggy import PluginManager as PluginManager_
+
+
+class AHook:
+    def __init__(self, pm: PluginManager_) -> None:
+        self.pm = pm
+
+    def __getattr__(self, name: str) -> Callable[..., Coroutine[Any, Any, list]]:
+        async def call(*args: Any, **kwargs: Any) -> list:
+            hook: HookCaller = getattr(self.pm.hook, name)
+            coros: list[asyncio.Future] = hook(*args, **kwargs)
+            return await asyncio.gather(*coros)
+
+        return call

src/apluggy/_wrap/awith.py

@@ -0,0 +1,53 @@
+import contextlib
+from collections.abc import AsyncIterator, Callable
+from typing import Any, AsyncContextManager
+
+from pluggy import HookCaller
+from pluggy import PluginManager as PluginManager_
+
+
+class AWith:
+    def __init__(self, pm: PluginManager_) -> None:
+        self.pm = pm
+
+    def __getattr__(self, name: str) -> Callable[..., AsyncContextManager]:
+        hook: HookCaller = getattr(self.pm.hook, name)
+        call = _Call(hook)
+        return call
+
+
+def _Call(
+    hook: Callable[..., list[AsyncContextManager]]
+) -> Callable[..., AsyncContextManager]:
+    @contextlib.asynccontextmanager
+    async def call(*args: Any, **kwargs: Any) -> AsyncIterator[list]:
+        ctxs = hook(*args, **kwargs)
+        async with contextlib.AsyncExitStack() as stack:
+            yields = [await stack.enter_async_context(ctx) for ctx in ctxs]
+
+            # TODO: Consider entering the contexts asynchronously as in the
+            # following commented out code.
+
+            # yields = await asyncio.gather(
+            #     *[stack.enter_async_context(ctx) for ctx in ctxs]
+            # )
+
+            yield yields
+
+            # TODO: The following commented out code is an attempt to support
+            # `asend()` through the `gen` attribute. It only works for
+            # simple cases. It doesn't work with starlette.lifespan().
+            # When starlette is shutting down, an exception is raised
+            # `RuntimeError: generator didn't stop after athrow()`.
+
+            # stop = False
+            # while not stop:
+            #     sent = yield yields
+            #     try:
+            #         yields = await asyncio.gather(
+            #             *[ctx.gen.asend(sent) for ctx in ctxs]
+            #         )
+            #     except StopAsyncIteration:
+            #         stop = True
+
+    return call