Add new sync server to separate file.

Author: ayjayt

src/py/kaleido/__init__.py

@@ -6,18 +6,30 @@
 
 from __future__ import annotations
 
-import asyncio
-import queue
-from threading import Thread
-from typing import TYPE_CHECKING, NamedTuple
-
 from choreographer.cli import get_chrome, get_chrome_sync
 
+from . import _sync_server
 from ._page_generator import PageGenerator
 from .kaleido import Kaleido
 
-if TYPE_CHECKING:
-    from typing import Any
+_global_server = _sync_server.GlobalKaleidoServer()
+
+
+def start_kaleido_sync_server(*args, **kwargs):
+    """
+    Start a kaleido server which will process all sync generation requests.
+
+    Only one server can be started at a time.
+
+    This wrapper function takes the exact same arguments as kaleido.Kaleido().
+    """
+    _global_server.open(*args, **kwargs)
+
+
+def stop_kaleido_sync_server():
+    """Stop the kaleido server. It can be restarted."""
+    _global_server.close()
+
 
 __all__ = [
     "Kaleido",
@@ -26,6 +38,8 @@
     "calc_fig_sync",
     "get_chrome",
     "get_chrome_sync",
+    "start_kaleido_sync_server",
+    "stop_kaleido_sync_server",
     "write_fig",
     "write_fig_from_object",
     "write_fig_from_object_sync",
@@ -126,36 +140,25 @@ async def write_fig_from_object(
         )
 
 
-def _async_thread_run(func, args: tuple[Any, ...], kwargs: dict):
-    q: queue.Queue[Any] = queue.Queue(maxsize=1)
-
-    def run(func, q, *args, **kwargs):
-        # func is a closure
-        try:
-            q.put(asyncio.run(func(*args, **kwargs)))
-        except BaseException as e:  # noqa: BLE001
-            q.put(e)
-
-    t = Thread(target=run, args=(func, q, *args), kwargs=kwargs)
-    t.start()
-    t.join()
-    res = q.get()
-    if isinstance(res, BaseException):
-        raise res
-    else:
-        return res
-
-
 def calc_fig_sync(*args, **kwargs):
     """Call `calc_fig` but blocking."""
-    return _async_thread_run(calc_fig, args=args, kwargs=kwargs)
+    if _global_server.is_running():
+        return _global_server.call_function("calc_fig", *args, **kwargs)
+    else:
+        return _sync_server.oneshot_async_run(calc_fig, args=args, kwargs=kwargs)
 
 
 def write_fig_sync(*args, **kwargs):
     """Call `write_fig` but blocking."""
-    _async_thread_run(write_fig, args=args, kwargs=kwargs)
+    if _global_server.is_running():
+        _global_server.call_function("write_fig", *args, **kwargs)
+    else:
+        _sync_server.oneshot_async_run(write_fig, args=args, kwargs=kwargs)
 
 
 def write_fig_from_object_sync(*args, **kwargs):
     """Call `write_fig_from_object` but blocking."""
-    _async_thread_run(write_fig_from_object, args=args, kwargs=kwargs)
+    if _global_server.is_running():
+        _global_server.call_function("write_fig_from_object", *args, **kwargs)
+    else:
+        _sync_server.oneshot_async_run(write_fig_from_object, args=args, kwargs=kwargs)

src/py/kaleido/_sync_server.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import asyncio
+import warnings
+from queue import Queue
+from threading import Thread
+from typing import TYPE_CHECKING, NamedTuple
+
+from .kaleido import Kaleido
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+class Task(NamedTuple):
+    fn: str
+    args: Any
+    kwargs: Any
+
+
+class _BadFunctionName(BaseException):
+    """For use when programmed poorly."""
+
+
+class GlobalKaleidoServer:
+    _instance = None
+
+    async def _server(self, *args, **kwargs):
+        async with Kaleido(*args, **kwargs) as k:  # multiple processor? Enable GPU?
+            while True:
+                task = self._task_queue.get()  # thread dies if main thread dies
+                if task is None:
+                    self._task_queue.task_done()
+                    return
+                if not hasattr(k, task.fn):
+                    raise _BadFunctionName(f"Kaleido has no attribute {task.fn}")
+                try:
+                    self._return_queue.put(
+                        await getattr(k, task.fn)(*task.args, **task.kwargs),
+                    )
+                except Exception as e:  # noqa: BLE001
+                    self._return_queue.put(e)
+
+                self._task_queue.task_done()
+
+    def __new__(cls):
+        # Create the singleton on first instantiation
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False  # noqa: SLF001
+        return cls._instance
+
+    def is_running(self):
+        return self._initialized
+
+    def open(self, *args, **kwargs):
+        """Initialize the singleton with three values."""
+        if self.is_running():
+            warnings.warn(
+                "Server already open.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return
+        coroutine = self._server(*args, **kwargs)
+        self._thread: Thread = Thread(target=asyncio.run, args=(coroutine,))
+        self._task_queue: Queue[Task | None] = Queue()
+        self._return_queue: Queue[Any] = Queue()
+        self._thread.start()
+        self._initialized = True
+
+    def close(self):
+        """Reset the singleton back to an uninitialized state."""
+        if not self.is_running():
+            warnings.warn(
+                "Server already closed.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return
+        self._task_queue.put(None)
+        self._thread.join()
+        del self._thread
+        del self._task_queue
+        del self._return_queue
+        self._initialized = False
+
+    def call_function(self, cmd: str, *args, **kwargs):
+        if not self._is_running():
+            raise RuntimeError("Can't call function on stopped server.")
+        if kwargs.pop("kopts"):
+            warnings.warn(
+                "The kopts argument is ignored if using a server.",
+                UserWarning,
+                stacklevel=3,
+            )
+        self._task_queue.put(Task(cmd, args, kwargs))
+        self._task_queue.join()
+        res = self._return_queue.get()
+        if isinstance(res, BaseException):
+            raise res
+        else:
+            return res
+
+
+def oneshot_async_run(func, args: tuple[Any, ...], kwargs: dict):
+    q: Queue[Any] = Queue(maxsize=1)
+
+    def run(func, q, *args, **kwargs):
+        # func is a closure
+        try:
+            q.put(asyncio.run(func(*args, **kwargs)))
+        except BaseException as e:  # noqa: BLE001
+            q.put(e)
+
+    t = Thread(target=run, args=(func, q, *args), kwargs=kwargs)
+    t.start()
+    t.join()
+    res = q.get()
+    if isinstance(res, BaseException):
+        raise res
+    else:
+        return res