浏览代码

introduce run.cpu_bound and run.io_bound

Falko Schindler 1 年之前
父节点
当前提交
291a74435f
共有 3 个文件被更改,包括 46 次插入2 次删除
  1. 4 1
      nicegui/__init__.py
  2. 3 1
      nicegui/nicegui.py
  3. 39 0
      nicegui/run_executor.py

+ 4 - 1
nicegui/__init__.py

@@ -1,4 +1,6 @@
-from . import elements, globals, ui  # pylint: disable=redefined-builtin
+from . import ui  # pylint: disable=redefined-builtin
+from . import elements, globals  # pylint: disable=redefined-builtin
+from . import run_executor as run
 from .api_router import APIRouter
 from .client import Client
 from .nicegui import app
@@ -11,6 +13,7 @@ __all__ = [
     'Client',
     'elements',
     'globals',
+    'run',
     'Tailwind',
     'ui',
     '__version__',

+ 3 - 1
nicegui/nicegui.py

@@ -11,7 +11,8 @@ from fastapi.responses import FileResponse, Response
 from fastapi.staticfiles import StaticFiles
 from fastapi_socketio import SocketManager
 
-from . import background_tasks, binding, favicon, globals, json, outbox, welcome  # pylint: disable=redefined-builtin
+from . import (background_tasks, binding, favicon, globals, json, outbox,  # pylint: disable=redefined-builtin
+               run_executor, welcome)
 from .app import App
 from .client import Client
 from .dependencies import js_components, libraries
@@ -109,6 +110,7 @@ async def handle_shutdown() -> None:
     with globals.index_client:
         for t in globals.shutdown_handlers:
             safe_invoke(t)
+    run_executor.tear_down()
     globals.state = globals.State.STOPPED
     if globals.air:
         await globals.air.disconnect()

+ 39 - 0
nicegui/run_executor.py

@@ -0,0 +1,39 @@
+import asyncio
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from functools import partial
+from typing import Any, Callable
+
+from . import globals, helpers  # pylint: disable=redefined-builtin
+
+process_pool = ProcessPoolExecutor()
+thread_pool = ThreadPoolExecutor()
+
+
+async def _run(executor: Any, callback: Callable, *args: Any, **kwargs: Any):
+    if globals.state == globals.State.STOPPING:
+        return
+    try:
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(executor, partial(callback, *args, **kwargs))
+    except RuntimeError as e:
+        if 'cannot schedule new futures after shutdown' not in str(e):
+            raise
+    except asyncio.exceptions.CancelledError:
+        pass
+
+
+async def cpu_bound(callback: Callable, *args: Any, **kwargs: Any):
+    _run(process_pool, callback, *args, **kwargs)
+
+
+async def io_bound(callback: Callable, *args: Any, **kwargs: Any):
+    _run(thread_pool, callback, *args, **kwargs)
+
+
+def tear_down() -> None:
+    if helpers.is_pytest():
+        return
+    for p in process_pool._processes.values():  # pylint: disable=protected-access
+        p.kill()
+    process_pool.shutdown(wait=True, cancel_futures=True)
+    thread_pool.shutdown(wait=False, cancel_futures=True)