diff --git a/bec_widgets/utils/bec_connector.py b/bec_widgets/utils/bec_connector.py index 91d53e33..81ef1fd4 100644 --- a/bec_widgets/utils/bec_connector.py +++ b/bec_widgets/utils/bec_connector.py @@ -6,7 +6,8 @@ import time import traceback import uuid from datetime import datetime -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Final, Optional +from weakref import WeakValueDictionary from bec_lib.logger import bec_logger from bec_lib.utils.import_utils import lazy_import_from @@ -29,6 +30,12 @@ else: logger = bec_logger.logger +class _NextSentinel(object): ... + + +NEXT_SENTINEL: Final[_NextSentinel] = _NextSentinel() + + class ConnectionConfig(BaseModel): """Configuration for BECConnector mixin class""" @@ -77,7 +84,8 @@ class BECConnector: """Connection mixin class to handle BEC client and device manager""" USER_ACCESS = ["_config_dict", "_get_all_rpc", "_rpc_id"] - EXIT_HANDLERS = {} + EXIT_HANDLERS: WeakValueDictionary[int, Callable[[],]] = WeakValueDictionary() + _exit_handler: Callable[[],] | None = None widget_removed = Signal() name_established = Signal(str) @@ -122,7 +130,7 @@ class BECConnector: self.client = self.bec_dispatcher.client if client is None else client self.rpc_register = RPCRegister() - if not self.client in BECConnector.EXIT_HANDLERS: + if not 0 in BECConnector.EXIT_HANDLERS.keys(): # register function to clean connections at exit; # the function depends on BECClient, and BECDispatcher @SafeSlot() @@ -143,8 +151,9 @@ class BECConnector: logger.info("Shutting down BEC Client", repr(client)) client.shutdown() - BECConnector.EXIT_HANDLERS[self.client] = terminate - QApplication.instance().aboutToQuit.connect(terminate) + BECConnector._exit_handler = terminate # type: ignore # keep a strong reference to the final cleanup + BECConnector._add_exit_handler(terminate, 0) + QApplication.instance().aboutToQuit.connect(self._run_exit_handlers) if config: self.config = config @@ -187,6 +196,42 @@ class BECConnector: QTimer.singleShot(0, self._update_object_name) + @classmethod + def _add_exit_handler(cls, handler: Callable, priority: int): + """Private to allow use of priority 0""" + cls.EXIT_HANDLERS[priority] = handler + + @classmethod + def add_exit_handler(cls, handler: Callable, priority: int | _NextSentinel = NEXT_SENTINEL): + """Add a handler to be called on the cleanup of the BEC Connector. Handlers are called in reverse order of their + priority - i.e. a higher number is higher priority. The BEC Connector's own cleanup will always be run last. + """ + existing_priorities = set(cls.EXIT_HANDLERS.keys()) + priority_modified = False + if isinstance(priority, _NextSentinel): + priority = max(existing_priorities) + 1 + if priority < 1: + raise ValueError( + "Please use a priority greater than 1! Priority 0 is reserved for system cleanup." + ) + if priority in cls.EXIT_HANDLERS.keys(): + priority_modified = True + logger.warning(f"Priority {priority} already in use - using the next available:") + while priority in cls.EXIT_HANDLERS.keys(): + priority += 1 + if priority_modified: + logger.warning(f"Assigned priority {priority} for {handler}.") + cls._add_exit_handler(handler, priority) + + @SafeSlot() + def _run_exit_handlers(self): + """Run all exit handlers from highest to lowest priority. Should be connected to AboutToQuit once and only once.""" + handlers = reversed( + list(handler for _, handler in sorted(BECConnector.EXIT_HANDLERS.items())) + ) + for handler in handlers: + handler() + @property def parent_id(self) -> str | None: try: diff --git a/tests/unit_tests/test_bec_connector.py b/tests/unit_tests/test_bec_connector.py index 62b88ca7..89ed4c07 100644 --- a/tests/unit_tests/test_bec_connector.py +++ b/tests/unit_tests/test_bec_connector.py @@ -1,5 +1,6 @@ # pylint: disable = no-name-in-module,missing-class-docstring, missing-module-docstring import time +from unittest.mock import MagicMock import pytest from qtpy.QtCore import QObject @@ -131,3 +132,20 @@ def test_bec_connector_change_object_name(bec_connector): # Verify that the object with the previous name is no longer registered all_objects = bec_connector.rpc_register.list_all_connections().values() assert not any(obj.objectName() == previous_name for obj in all_objects) + + +def test_bec_connector_terminate_run_on_about_to_quit(bec_connector): + assert BECConnector.EXIT_HANDLERS.get(0) is not None + terminate_mock = MagicMock() + bec_connector.__class__.EXIT_HANDLERS[0] = terminate_mock + QApplication.instance().aboutToQuit.emit() + terminate_mock.assert_called_once() + + +def test_bec_connector_terminate_run_once_and_only_once(bec_connector): + terminate_mock = MagicMock() + bec_connector.__class__.EXIT_HANDLERS[0] = terminate_mock + conn_2 = BECConnectorQObject(client=mocked_client) + conn_3 = BECConnectorQObject(client=mocked_client) + QApplication.instance().aboutToQuit.emit() + terminate_mock.assert_called_once()