diff --git a/src/py/kaleido/__init__.py b/src/py/kaleido/__init__.py index 5cfd03bd..d9961019 100644 --- a/src/py/kaleido/__init__.py +++ b/src/py/kaleido/__init__.py @@ -12,25 +12,6 @@ from ._page_generator import PageGenerator from .kaleido import Kaleido -_global_server = _sync_server.GlobalKaleidoServer() - - -def start_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_sync_server(): - """Stop the kaleido server. It can be restarted.""" - _global_server.close() - - __all__ = [ "Kaleido", "PageGenerator", @@ -46,6 +27,40 @@ def stop_sync_server(): "write_fig_sync", ] +_global_server = _sync_server.GlobalKaleidoServer() + + +def start_sync_server(*args, silence_warnings=False, **kwargs): + """ + Start a kaleido server which will process all sync generation requests. + + The kaleido server is a singleton, so it can't be opened twice. This + function will warn you if the server is already running. + + This wrapper function takes the exact same arguments as kaleido.Kaleido(), + except one extra, `silence_warnings`. + + Args: + *args: all arguments `Kaleido()` would take. + silence_warnings: (bool, default False): If True, don't emit warning if + starting an already started server. + **kwargs: all keyword arguments `Kaleido()` would take. + + """ + _global_server.open(*args, silence_warnings=silence_warnings, **kwargs) + + +def stop_sync_server(*, silence_warnings=False): + """ + Stop the kaleido server. It can be restarted. Warns if not started. + + Args: + silence_warnings: (bool, default False): If True, don't emit warning if + stopping a server that's not running. + + """ + _global_server.close(silence_warnings=silence_warnings) + async def calc_fig( fig, diff --git a/src/py/kaleido/_sync_server.py b/src/py/kaleido/_sync_server.py index b765b433..714fcce2 100644 --- a/src/py/kaleido/_sync_server.py +++ b/src/py/kaleido/_sync_server.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import atexit import warnings +from functools import partial from queue import Queue from threading import Thread from typing import TYPE_CHECKING, NamedTuple @@ -53,30 +55,38 @@ def __new__(cls): def is_running(self): return self._initialized - def open(self, *args: Any, **kwargs: Any) -> None: + def open(self, *args: Any, silence_warnings=False, **kwargs: Any) -> None: """Initialize the singleton with three values.""" if self.is_running(): - warnings.warn( - "Server already open.", - RuntimeWarning, - stacklevel=2, - ) + if not silence_warnings: + warnings.warn( + "Server already open.", + RuntimeWarning, + stacklevel=2, + ) return coroutine = self._server(*args, **kwargs) - self._thread: Thread = Thread(target=asyncio.run, args=(coroutine,)) + self._thread: Thread = Thread( + target=asyncio.run, + args=(coroutine,), + daemon=True, + ) self._task_queue: Queue[Task | None] = Queue() self._return_queue: Queue[Any] = Queue() self._thread.start() self._initialized = True + close = partial(self.close, silence_warnings=True) + atexit.register(close) - def close(self): + def close(self, *, silence_warnings=False): """Reset the singleton back to an uninitialized state.""" if not self.is_running(): - warnings.warn( - "Server already closed.", - RuntimeWarning, - stacklevel=2, - ) + if not silence_warnings: + warnings.warn( + "Server already closed.", + RuntimeWarning, + stacklevel=2, + ) return self._task_queue.put(None) self._thread.join()