Feat: Added support for running async context manager with sync_worker

Author: eadwinCode

ellar/app/context.py

@@ -16,6 +16,7 @@
 app_context_var: ContextVar[
     t.Optional[t.Union["ApplicationContext", t.Any]]
 ] = ContextVar("ellar.app.context")
+app_context_var.set(empty)
 
 
 class ApplicationContext:

ellar/common/utils/importer.py

@@ -96,10 +96,10 @@ def get_main_directory_by_stack(
     if forced_path_to_string.startswith("__main__") or forced_path_to_string.startswith(
         "/__main__"
     ):
-        __main__, others = forced_path_to_string.replace("/", " ").split("__main__")
+        __main__, *others = forced_path_to_string.replace("/", " ").split("__main__")
         __parent__ = False
 
-        if "__parent__" in others:
+        if "__parent__" in others[0]:
             __parent__ = True
 
         if not from_dir:
@@ -110,12 +110,13 @@ def get_main_directory_by_stack(
             __main__parent = Path(from_dir).resolve()
 
         if __parent__:
-            parent_split = others.split("__parent__")
-            for item in parent_split:
+            others = others[0].split("__parent__")
+            for item in list(others):
                 if item == " ":
                     __main__parent = __main__parent.parent
-                else:
-                    return os.path.join(__main__parent, item.strip())
+                    others.remove(item)
 
-        return os.path.join(str(__main__parent), others.strip())
+        return os.path.join(
+            str(__main__parent), *[i.strip() for i in others[0].split(" ")]
+        )
     return path

ellar/threading/sync_worker.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import enum
 import inspect
 import logging
@@ -36,7 +37,9 @@ def __bool__(self) -> bool:  # pragma: no cover
 class _SyncWorkerThread(threading.Thread):
     work_queue: queue.Queue[
         t.Union[
-            t.Tuple[t.Union[t.AsyncIterator, t.Coroutine], Context],
+            t.Tuple[
+                t.Union[t.AsyncIterator, t.Coroutine, t.AsyncContextManager], Context
+            ],
             _Sentinel,
         ]
     ]
@@ -72,6 +75,11 @@ def run(self) -> None:
                 coro, ctx = item
                 if inspect.isasyncgen(coro):
                     ctx.run(loop.run_until_complete, self.agen_wrapper(coro))  # type: ignore[arg-type]
+                elif isinstance(coro, t.AsyncContextManager):
+                    ctx.run(
+                        loop.run_until_complete,
+                        self.async_context_manager_wrapper(coro),
+                    )
                 else:
                     try:
                         # 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
                     if item is sentinel:
                         break
                     if isinstance(item, Exception):
-                        self.work_queue.put(sentinel)  # initial loop closing
                         raise item
                     yield item
                 finally:
                     self.stream_block.set()
                     self.stream_queue.task_done()
         finally:
             del ctx
+            self.work_queue.put(sentinel)  # initial loop closing
+
+    def _update_context(self, context: Context) -> None:
+        for var, value in context.items():
+            var.set(value)
+
+    @contextlib.contextmanager
+    def execute_async_context_generator(
+        self, async_context_manager: t.AsyncContextManager, context_update: bool = True
+    ) -> t.Generator:
+        ctx = copy_context()  # preserve context for the worker thread
+
+        try:
+            self.work_queue.put((async_context_manager, ctx))
+            item, updated_ctx = self.stream_queue.get()  # type:ignore[misc]
+
+            try:
+                if isinstance(item, Exception):
+                    raise item
+
+                if context_update:
+                    self._update_context(updated_ctx)
+
+                yield item
+            finally:
+                if updated_ctx:
+                    del updated_ctx
+                    self._update_context(ctx)
+
+                self.stream_block.set()
+                self.stream_queue.task_done()
+        finally:
+            del ctx
+            self.work_queue.put(sentinel)  # initial loop closing
+
+    async def async_context_manager_wrapper(self, agen: t.AsyncContextManager) -> None:
+        try:
+            async with agen as s:
+                self.stream_block.clear()
+                self.stream_queue.put((s, copy_context()))
+                # flow-control the generator.
+                self.stream_block.wait()
+        except Exception as e:
+            self.stream_queue.put((e, None))
 
     def interrupt_generator(self) -> None:
         self.agen_shutdown = True
@@ -169,3 +220,19 @@ def execute_async_gen_with_sync_worker(
 
     _worker_thread.work_queue.put(sentinel)
     _worker_thread.join()
+
+
+@contextlib.contextmanager  # type:ignore[arg-type]
+def execute_async_context_manager_with_sync_worker(  # type:ignore[misc]
+    async_gen: t.AsyncContextManager, context_update: bool = True
+) -> t.ContextManager:
+    _worker_thread = _SyncWorkerThread()
+    _worker_thread.start()
+
+    with _worker_thread.execute_async_context_generator(
+        async_gen, context_update
+    ) as item:
+        yield item
+
+    _worker_thread.work_queue.put(sentinel)
+    _worker_thread.join()

tests/test_thread_worker/test_sync_worker.py

@@ -1,6 +1,9 @@
+import contextlib
+
 import pytest
 from ellar.threading.sync_worker import (
     _SyncWorkerThread,
+    execute_async_context_manager_with_sync_worker,
     execute_async_gen_with_sync_worker,
     execute_coroutine_with_sync_worker,
     sentinel,
@@ -22,6 +25,13 @@ async def async_gen(after=None):
         yield i
 
 
+@contextlib.asynccontextmanager
+async def async_context_manager(with_exception=False):
+    if with_exception:
+        raise RuntimeError("Context Manager Raised an Exception")
+    yield 10
+
+
 async def test_run_with_sync_worker_runs_async_function_synchronously(anyio_backend):
     res = execute_coroutine_with_sync_worker(coroutine_function())
     assert res == "Coroutine Function"
@@ -71,3 +81,14 @@ async def test_sync_worker_interrupt_function_works(anyio_backend):
     assert res == [0, 1, 2, 3, 4, 5, 6]
     worker.work_queue.put(sentinel)
     worker.join()
+
+
+async def test_sync_worker_runs_async_context_manager(anyio_backend):
+    with execute_async_context_manager_with_sync_worker(async_context_manager()) as ctx:
+        assert ctx == 10
+
+    with pytest.raises(RuntimeError, match="Context Manager Raised an Exception"):
+        with execute_async_context_manager_with_sync_worker(
+            async_context_manager(with_exception=True)
+        ) as ctx:
+            pass

tests/test_utils/test_importer.py

@@ -30,3 +30,13 @@ def test_get_main_directory_by_stack_works_with_a_path_reference():
     result = os.listdir(directory)
 
     assert result == ["test.css"]
+
+
+def test_get_main_directory_by_stack_works_path_chaining():
+    path = "__main__/dumbs/default"
+    directory = get_main_directory_by_stack(path, stack_level=1)
+    assert "/test_utils/dumbs/default" in directory
+
+    path = "__main__/__parent__/dumbs/default"
+    directory = get_main_directory_by_stack(path, stack_level=1)
+    assert "/tests/dumbs/default" in directory