Skip to content

Commit 99be2a2

Browse files
add utilities to deal with async in pluggy
1 parent 02095a3 commit 99be2a2

12 files changed

+651
-19
lines changed

.pre-commit-config.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,4 +48,4 @@ repos:
4848
- id: mypy
4949
files: ^(src/|testing/)
5050
args: []
51-
additional_dependencies: [pytest]
51+
additional_dependencies: [pytest, types-greenlet]

pyproject.toml

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,14 @@ readme = {file = "README.rst", content-type = "text/x-rst"}
3535
requires-python = ">=3.9"
3636

3737
dynamic = ["version"]
38+
39+
[project.optional-dependencies]
40+
async = ["greenlet"]
41+
3842
[dependency-groups]
3943
dev = ["pre-commit", "tox", "mypy", "ruff"]
40-
testing = ["pytest", "pytest-benchmark", "coverage"]
44+
testing = ["pytest", "pytest-benchmark", "coverage", "greenlet", "types-greenlet"]
45+
4146

4247

4348
[tool.setuptools]
@@ -46,6 +51,7 @@ package-dir = {""="src"}
4651
package-data = {"pluggy" = ["py.typed"]}
4752

4853

54+
4955
[tool.ruff.lint]
5056
extend-select = [
5157
"I", # isort

src/pluggy/_async.py

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
"""
2+
Async support for pluggy using greenlets.
3+
4+
This module provides async functionality for pluggy, allowing hook implementations
5+
to return awaitable objects that are automatically awaited when running in an
6+
async context.
7+
"""
8+
9+
from __future__ import annotations
10+
11+
from collections.abc import AsyncGenerator
12+
from collections.abc import Awaitable
13+
from collections.abc import Callable
14+
from collections.abc import Generator
15+
from typing import Any
16+
from typing import TYPE_CHECKING
17+
from typing import TypeVar
18+
19+
20+
_T = TypeVar("_T")
21+
_Y = TypeVar("_Y")
22+
_S = TypeVar("_S")
23+
_R = TypeVar("_R")
24+
25+
if TYPE_CHECKING:
26+
import greenlet
27+
28+
29+
def make_greenlet(func: Callable[..., Any]) -> greenlet.greenlet:
30+
"""indirection to defer import"""
31+
import greenlet
32+
33+
return greenlet.greenlet(func)
34+
35+
36+
class Submitter:
37+
# practice we expect te root greenlet to be the key submitter
38+
_active_submitter: greenlet.greenlet | None
39+
40+
def __init__(self) -> None:
41+
self._active_submitter = None
42+
43+
def __repr__(self) -> str:
44+
return f"<Submitter active={self._active_submitter is not None}>"
45+
46+
def maybe_submit(self, coro: Awaitable[_T]) -> _T | Awaitable[_T]:
47+
"""await an awaitable if active, else return it
48+
49+
this enables backward compatibility for datasette
50+
and https://simonwillison.net/2020/Sep/2/await-me-maybe/
51+
"""
52+
active = self._active_submitter
53+
if active is not None:
54+
# We're in a greenlet context, switch with the awaitable
55+
# The parent will await it and switch back with the result
56+
res: _T = active.switch(coro)
57+
return res
58+
else:
59+
return coro
60+
61+
def require_await(self, coro: Awaitable[_T]) -> _T:
62+
"""await an awaitable, raising an error if not in async context
63+
64+
this is for cases where async context is required
65+
"""
66+
active = self._active_submitter
67+
if active is not None:
68+
# Switch to the active submitter greenlet with the awaitable
69+
# The active submitter will await it and switch back with the result
70+
res: _T = active.switch(coro)
71+
return res
72+
else:
73+
raise RuntimeError("require_await called outside of async context")
74+
75+
async def run(self, sync_func: Callable[[], _T]) -> _T:
76+
"""Run a synchronous function with async support."""
77+
try:
78+
import greenlet
79+
except ImportError:
80+
raise RuntimeError("greenlet is required for async support")
81+
82+
if self._active_submitter is not None:
83+
raise RuntimeError("Submitter is already active")
84+
85+
# Set the current greenlet as the main async context
86+
main_greenlet = greenlet.getcurrent()
87+
result: _T | None = None
88+
exception: BaseException | None = None
89+
90+
def greenlet_func() -> None:
91+
nonlocal result, exception
92+
try:
93+
result = sync_func()
94+
except BaseException as e:
95+
exception = e
96+
97+
# Create the worker greenlet
98+
worker_greenlet = greenlet.greenlet(greenlet_func)
99+
# Set the active submitter to the main greenlet so maybe_submit can switch back
100+
self._active_submitter = main_greenlet
101+
102+
try:
103+
# Switch to the worker greenlet and handle any awaitables it passes back
104+
awaitable = worker_greenlet.switch()
105+
while awaitable is not None:
106+
# Await the awaitable and send the result back to the greenlet
107+
awaited_result = await awaitable
108+
awaitable = worker_greenlet.switch(awaited_result)
109+
except Exception as e:
110+
# If something goes wrong, try to send the exception to the greenlet
111+
try:
112+
worker_greenlet.throw(e)
113+
except BaseException as inner_e:
114+
exception = inner_e
115+
finally:
116+
self._active_submitter = None
117+
118+
if exception is not None:
119+
raise exception
120+
if result is None:
121+
raise RuntimeError("Function completed without setting result")
122+
return result
123+
124+
125+
def async_generator_to_sync(
126+
async_gen: AsyncGenerator[_Y, _S], submitter: Submitter
127+
) -> Generator[_Y, _S, None]:
128+
"""Convert an async generator to a sync generator using a Submitter.
129+
130+
This helper allows wrapper implementations to use async generators while
131+
maintaining compatibility with the sync generator interface expected by
132+
the hook system.
133+
134+
Args:
135+
async_gen: The async generator to convert
136+
submitter: The Submitter to use for awaiting async operations
137+
138+
Yields:
139+
Values from the async generator
140+
141+
Returns:
142+
None (async generators don't return values)
143+
144+
Example:
145+
async def my_async_wrapper():
146+
yield # Setup phase
147+
result = await some_async_operation()
148+
149+
# In a wrapper hook implementation:
150+
def my_wrapper_hook():
151+
async_gen = my_async_wrapper()
152+
gen = async_generator_to_sync(async_gen, submitter)
153+
try:
154+
while True:
155+
value = next(gen)
156+
yield value
157+
except StopIteration:
158+
return
159+
"""
160+
try:
161+
# Start the async generator
162+
value = submitter.require_await(async_gen.__anext__())
163+
164+
while True:
165+
try:
166+
# Yield the value and get the sent value
167+
sent_value = yield value
168+
169+
# Send the value to the async generator and get the next value
170+
try:
171+
value = submitter.require_await(async_gen.asend(sent_value))
172+
except StopAsyncIteration:
173+
# Async generator completed
174+
return
175+
176+
except GeneratorExit:
177+
# Generator is being closed, close the async generator
178+
try:
179+
submitter.require_await(async_gen.aclose())
180+
except StopAsyncIteration:
181+
pass
182+
raise
183+
184+
except BaseException as exc:
185+
# Exception was thrown into the generator,
186+
# throw it into the async generator
187+
try:
188+
value = submitter.require_await(async_gen.athrow(exc))
189+
except StopAsyncIteration:
190+
# Async generator completed
191+
return
192+
except StopIteration as sync_stop_exc:
193+
# Re-raise StopIteration as it was passed through
194+
raise sync_stop_exc
195+
196+
except StopAsyncIteration:
197+
# Async generator completed normally
198+
return

src/pluggy/_callers.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,18 @@
44

55
from __future__ import annotations
66

7+
from collections.abc import Awaitable
78
from collections.abc import Generator
89
from collections.abc import Mapping
910
from collections.abc import Sequence
1011
from typing import cast
1112
from typing import NoReturn
13+
from typing import TYPE_CHECKING
1214
import warnings
1315

16+
17+
if TYPE_CHECKING:
18+
from ._async import Submitter
1419
from ._hook_callers import HookImpl
1520
from ._hook_callers import WrapperImpl
1621
from ._result import Result
@@ -79,6 +84,7 @@ def _multicall(
7984
wrapper_impls: Sequence[WrapperImpl],
8085
caller_kwargs: Mapping[str, object],
8186
firstresult: bool,
87+
async_submitter: Submitter,
8288
) -> object | list[object]:
8389
"""Execute a call into multiple python functions/methods and return the
8490
result(s).
@@ -108,6 +114,9 @@ def _multicall(
108114
args = normal_impl._get_call_args(caller_kwargs)
109115
res = normal_impl.function(*args)
110116
if res is not None:
117+
# Handle awaitable results using maybe_submit
118+
if isinstance(res, Awaitable):
119+
res = async_submitter.maybe_submit(res)
111120
results.append(res)
112121
if firstresult: # halt further impl calls
113122
break

0 commit comments

Comments
 (0)