Skip to content

Commit d970e8e

Browse files
authored
Merge pull request #169 from python-ellar/sync_context_worker
Feat: Added support for running async context manager with sync_worker
2 parents f0d77a2 + 30a3a10 commit d970e8e

File tree

6 files changed

+110
-10
lines changed

6 files changed

+110
-10
lines changed

ellar/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
""" Ellar - Python ASGI web framework for building fast, efficient, and scalable RESTful APIs and server-side applications. """
22

3-
__version__ = "0.6.5"
3+
__version__ = "0.6.6"

ellar/app/context.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
app_context_var: ContextVar[
1717
t.Optional[t.Union["ApplicationContext", t.Any]]
1818
] = ContextVar("ellar.app.context")
19+
app_context_var.set(empty)
1920

2021

2122
class ApplicationContext:

ellar/common/utils/importer.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,10 +96,10 @@ def get_main_directory_by_stack(
9696
if forced_path_to_string.startswith("__main__") or forced_path_to_string.startswith(
9797
"/__main__"
9898
):
99-
__main__, others = forced_path_to_string.replace("/", " ").split("__main__")
99+
__main__, *others = forced_path_to_string.replace("/", " ").split("__main__")
100100
__parent__ = False
101101

102-
if "__parent__" in others:
102+
if "__parent__" in others[0]:
103103
__parent__ = True
104104

105105
if not from_dir:
@@ -110,12 +110,13 @@ def get_main_directory_by_stack(
110110
__main__parent = Path(from_dir).resolve()
111111

112112
if __parent__:
113-
parent_split = others.split("__parent__")
114-
for item in parent_split:
113+
others = others[0].split("__parent__")
114+
for item in list(others):
115115
if item == " ":
116116
__main__parent = __main__parent.parent
117-
else:
118-
return os.path.join(__main__parent, item.strip())
117+
others.remove(item)
119118

120-
return os.path.join(str(__main__parent), others.strip())
119+
return os.path.join(
120+
str(__main__parent), *[i.strip() for i in others[0].split(" ")]
121+
)
121122
return path

ellar/threading/sync_worker.py

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from __future__ import annotations
66

77
import asyncio
8+
import contextlib
89
import enum
910
import inspect
1011
import logging
@@ -36,7 +37,9 @@ def __bool__(self) -> bool: # pragma: no cover
3637
class _SyncWorkerThread(threading.Thread):
3738
work_queue: queue.Queue[
3839
t.Union[
39-
t.Tuple[t.Union[t.AsyncIterator, t.Coroutine], Context],
40+
t.Tuple[
41+
t.Union[t.AsyncIterator, t.Coroutine, t.AsyncContextManager], Context
42+
],
4043
_Sentinel,
4144
]
4245
]
@@ -72,6 +75,11 @@ def run(self) -> None:
7275
coro, ctx = item
7376
if inspect.isasyncgen(coro):
7477
ctx.run(loop.run_until_complete, self.agen_wrapper(coro)) # type: ignore[arg-type]
78+
elif isinstance(coro, t.AsyncContextManager):
79+
ctx.run(
80+
loop.run_until_complete,
81+
self.async_context_manager_wrapper(coro),
82+
)
7583
else:
7684
try:
7785
# FIXME: Once python/mypy#12756 is resolved, remove the type-ignore tag.
@@ -131,14 +139,57 @@ def execute_generator(self, async_gen: t.AsyncIterator[_Item]) -> t.Iterator[_It
131139
if item is sentinel:
132140
break
133141
if isinstance(item, Exception):
134-
self.work_queue.put(sentinel) # initial loop closing
135142
raise item
136143
yield item
137144
finally:
138145
self.stream_block.set()
139146
self.stream_queue.task_done()
140147
finally:
141148
del ctx
149+
self.work_queue.put(sentinel) # initial loop closing
150+
151+
def _update_context(self, context: Context) -> None:
152+
for var, value in context.items():
153+
var.set(value)
154+
155+
@contextlib.contextmanager
156+
def execute_async_context_generator(
157+
self, async_context_manager: t.AsyncContextManager, context_update: bool = True
158+
) -> t.Generator:
159+
ctx = copy_context() # preserve context for the worker thread
160+
161+
try:
162+
self.work_queue.put((async_context_manager, ctx))
163+
item, updated_ctx = self.stream_queue.get() # type:ignore[misc]
164+
165+
try:
166+
if isinstance(item, Exception):
167+
raise item
168+
169+
if context_update:
170+
self._update_context(updated_ctx)
171+
172+
yield item
173+
finally:
174+
if updated_ctx:
175+
del updated_ctx
176+
self._update_context(ctx)
177+
178+
self.stream_block.set()
179+
self.stream_queue.task_done()
180+
finally:
181+
del ctx
182+
self.work_queue.put(sentinel) # initial loop closing
183+
184+
async def async_context_manager_wrapper(self, agen: t.AsyncContextManager) -> None:
185+
try:
186+
async with agen as s:
187+
self.stream_block.clear()
188+
self.stream_queue.put((s, copy_context()))
189+
# flow-control the generator.
190+
self.stream_block.wait()
191+
except Exception as e:
192+
self.stream_queue.put((e, None))
142193

143194
def interrupt_generator(self) -> None:
144195
self.agen_shutdown = True
@@ -169,3 +220,19 @@ def execute_async_gen_with_sync_worker(
169220

170221
_worker_thread.work_queue.put(sentinel)
171222
_worker_thread.join()
223+
224+
225+
@contextlib.contextmanager # type:ignore[arg-type]
226+
def execute_async_context_manager_with_sync_worker( # type:ignore[misc]
227+
async_gen: t.AsyncContextManager, context_update: bool = True
228+
) -> t.ContextManager:
229+
_worker_thread = _SyncWorkerThread()
230+
_worker_thread.start()
231+
232+
with _worker_thread.execute_async_context_generator(
233+
async_gen, context_update
234+
) as item:
235+
yield item
236+
237+
_worker_thread.work_queue.put(sentinel)
238+
_worker_thread.join()

tests/test_thread_worker/test_sync_worker.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
import contextlib
2+
13
import pytest
24
from ellar.threading.sync_worker import (
35
_SyncWorkerThread,
6+
execute_async_context_manager_with_sync_worker,
47
execute_async_gen_with_sync_worker,
58
execute_coroutine_with_sync_worker,
69
sentinel,
@@ -22,6 +25,13 @@ async def async_gen(after=None):
2225
yield i
2326

2427

28+
@contextlib.asynccontextmanager
29+
async def async_context_manager(with_exception=False):
30+
if with_exception:
31+
raise RuntimeError("Context Manager Raised an Exception")
32+
yield 10
33+
34+
2535
async def test_run_with_sync_worker_runs_async_function_synchronously(anyio_backend):
2636
res = execute_coroutine_with_sync_worker(coroutine_function())
2737
assert res == "Coroutine Function"
@@ -71,3 +81,14 @@ async def test_sync_worker_interrupt_function_works(anyio_backend):
7181
assert res == [0, 1, 2, 3, 4, 5, 6]
7282
worker.work_queue.put(sentinel)
7383
worker.join()
84+
85+
86+
async def test_sync_worker_runs_async_context_manager(anyio_backend):
87+
with execute_async_context_manager_with_sync_worker(async_context_manager()) as ctx:
88+
assert ctx == 10
89+
90+
with pytest.raises(RuntimeError, match="Context Manager Raised an Exception"):
91+
with execute_async_context_manager_with_sync_worker(
92+
async_context_manager(with_exception=True)
93+
) as ctx:
94+
pass

tests/test_utils/test_importer.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,13 @@ def test_get_main_directory_by_stack_works_with_a_path_reference():
3030
result = os.listdir(directory)
3131

3232
assert result == ["test.css"]
33+
34+
35+
def test_get_main_directory_by_stack_works_path_chaining():
36+
path = "__main__/dumbs/default"
37+
directory = get_main_directory_by_stack(path, stack_level=1)
38+
assert "/test_utils/dumbs/default" in directory
39+
40+
path = "__main__/__parent__/dumbs/default"
41+
directory = get_main_directory_by_stack(path, stack_level=1)
42+
assert "/tests/dumbs/default" in directory

0 commit comments

Comments
 (0)