Callback decorator

Author: elacuesta

scrapy_playwright/_utils.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import logging
 import platform
 import threading
@@ -8,6 +9,7 @@
 from playwright.async_api import Error, Page, Request, Response
 from scrapy.http.headers import Headers
 from scrapy.settings import Settings
+from scrapy.utils.asyncgen import collect_asyncgen
 from scrapy.utils.python import to_unicode
 from twisted.internet.defer import Deferred
 from w3lib.encoding import html_body_declared_encoding, http_content_type_encoding
@@ -117,6 +119,7 @@ class _ThreadedLoopAdapter:
     @classmethod
     async def _handle_coro(cls, coro, future) -> None:
         try:
+            coro = collect_asyncgen(coro) if inspect.isasyncgen(coro) else coro
             future.set_result(await coro)
         except Exception as exc:
             future.set_exception(exc)

scrapy_playwright/utils.py

@@ -1,13 +1,18 @@
-import asyncio
 import functools
 import inspect
-from typing import Awaitable
+from typing import Awaitable, Callable
 
 from ._utils import _ThreadedLoopAdapter
 
 
-def ensure_future(coro: Awaitable) -> asyncio.Future:
-    """Wrap a coroutine in a Future assigned to the threaded event loop.
+async def _run_async_gen(asyncgen):
+    async for item in asyncgen:
+        yield item
+
+
+def use_threaded_loop(callback: Awaitable) -> Callable:
+    """Wrap a coroutine callback so that Playwright coroutines are executed in
+    the threaded event loop.
 
     On windows, Playwright runs in an event loop of its own in a separate thread.
     If Playwright coroutines are awaited directly, they are assigned to the main
@@ -17,26 +22,33 @@ def ensure_future(coro: Awaitable) -> asyncio.Future:
     Usage:
     ```
     from playwright.async_api import Page
-    from scrapy_playwright import ensure_future
+    from scrapy_playwright.utils import use_threaded_loop
 
+    @use_threaded_loop
     async def parse(self, response):
         page: Page = response.meta["playwright_page"]
-        await ensure_future(page.screenshot(path="example.png", full_page=True))
+        await page.screenshot(path="example.png", full_page=True)
     ```
     """
-    return _ThreadedLoopAdapter._ensure_future(coro)
 
-
-def use_threaded_loop(callback):
-    if not (inspect.iscoroutinefunction(callback) or inspect.isasyncgenfunction(callback)):
+    if not inspect.iscoroutinefunction(callback) and not inspect.isasyncgenfunction(callback):
         raise RuntimeError(
             f"Cannot decorate callback '{callback.__name__}' with 'use_threaded_loop':"
             " callback must be a coroutine function or an async generator"
         )
 
     @functools.wraps(callback)
-    async def wrapper(*args, **kwargs):
-        future: asyncio.Future = _ThreadedLoopAdapter._ensure_future(callback(*args, **kwargs))
+    async def async_func_wrapper(*args, **kwargs):
+        future = _ThreadedLoopAdapter._ensure_future(callback(*args, **kwargs))
         return await future
 
-    return wrapper
+    @functools.wraps(callback)
+    async def async_gen_wrapper(*args, **kwargs):
+        asyncgen = _run_async_gen(callback(*args, **kwargs))
+        future = _ThreadedLoopAdapter._ensure_future(asyncgen)
+        for item in await future:
+            yield item
+
+    if inspect.isasyncgenfunction(callback):
+        return async_gen_wrapper
+    return async_func_wrapper