moves state manager from DataService to Server

src/pydase/data_service/abstract_data_service.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Optional
 if TYPE_CHECKING:
     from pydase.data_service.callback_manager import CallbackManager
     from pydase.data_service.data_service import DataService
-    from pydase.data_service.state_manager import StateManager
     from pydase.data_service.task_manager import TaskManager
 
 
@@ -14,6 +13,5 @@ class AbstractDataService(ABC):
     __root__: DataService
     _task_manager: TaskManager
     _callback_manager: CallbackManager
-    _state_manager: Optional[StateManager]
     _autostart_tasks: dict[str, tuple[Any]]
     _filename: Optional[str]

src/pydase/data_service/data_service.py

@@ -43,7 +43,6 @@ class DataService(rpyc.Service, AbstractDataService):
         self._filename: Optional[str] = filename
         self._callback_manager: CallbackManager = CallbackManager(self)
         self._task_manager = TaskManager(self)
-        self._state_manager: Optional[StateManager] = None
 
         if not hasattr(self, "_autostart_tasks"):
             self._autostart_tasks = {}
@@ -56,10 +55,6 @@ class DataService(rpyc.Service, AbstractDataService):
         self.__check_instance_classes()
         self._initialised = True
 
-        if self._filename is not None:
-            self._state_manager = StateManager(self)
-            self._state_manager.load_state()
-
     def __setattr__(self, __name: str, __value: Any) -> None:
         # converting attributes that are not properties
         if not isinstance(getattr(type(self), __name, None), property):
@@ -135,11 +130,9 @@ class DataService(rpyc.Service, AbstractDataService):
         """
         Serialize the DataService instance and write it to a JSON file.
 
-        Args:
-            filename (str): The name of the file to write to.
         """
-        if self._state_manager is not None:
-            self._state_manager.save_state()
+        if hasattr(self, "_state_manager"):
+            getattr(self, "_state_manager").save_state()
 
     def load_DataService_from_JSON(self, json_dict: dict[str, Any]) -> None:
         # Traverse the serialized representation and set the attributes of the class

src/pydase/server/server.py

@@ -5,6 +5,7 @@ import signal
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
+from pathlib import Path
 from types import FrameType
 from typing import Any, Optional, Protocol, TypedDict
 
@@ -158,15 +159,13 @@ class Server:
         web_port: int = 8001,
         enable_rpc: bool = True,
         enable_web: bool = True,
+        filename: Optional[str | Path] = None,
         use_forking_server: bool = False,
         web_settings: dict[str, Any] = {},
         additional_servers: list[AdditionalServer] = [],
         **kwargs: Any,
     ) -> None:
         self._service = service
-        if self._service._state_manager is None:
-            self._service._state_manager = StateManager(self._service)
-            self._service._state_manager.load_state()
         self._host = host
         self._rpc_port = rpc_port
         self._web_port = web_port
@@ -191,6 +190,10 @@ class Server:
             "additional_servers": [],
             **kwargs,
         }
+        self._state_manager = StateManager(self._service, filename)
+        if self._service._filename is not None:
+            self._service._state_manager = self._state_manager
+        self._state_manager.load_state()
 
     def run(self) -> None:
         """
@@ -275,6 +278,7 @@ class Server:
             self._wapi: WebAPI = WebAPI(
                 service=self._service,
                 info=self._info,
+                state_manager=self._state_manager,
                 **self._kwargs,
             )
             web_server = uvicorn.Server(
@@ -326,9 +330,9 @@ class Server:
     async def shutdown(self) -> None:
         logger.info("Shutting down")
 
-        logger.info(f"Saving data to {self._service._filename}.")
-        if self._service._state_manager is not None:
-            self._service._state_manager.save_state()
+        logger.info(f"Saving data to {self._state_manager.filename}.")
+        if self._state_manager is not None:
+            self._state_manager.save_state()
 
         await self.__cancel_servers()
         await self.__cancel_tasks()

src/pydase/server/web_server.py

@@ -9,6 +9,7 @@ from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
 from pydase import DataService
+from pydase.data_service.state_manager import StateManager
 from pydase.version import __version__
 
 logger = logging.getLogger(__name__)
@@ -50,6 +51,7 @@ class WebAPI:
     def __init__(  # noqa: CFQ002
         self,
         service: DataService,
+        state_manager: StateManager,
         frontend: str | Path | None = None,
         css: str | Path | None = None,
         enable_CORS: bool = True,
@@ -58,6 +60,7 @@ class WebAPI:
         **kwargs: Any,
     ):
         self.service = service
+        self.state_manager = state_manager
         self.frontend = frontend
         self.css = css
         self.enable_CORS = enable_CORS
@@ -114,11 +117,7 @@ class WebAPI:
 
         @app.get("/service-properties")
         def service_properties() -> dict[str, Any]:
-            if self.service._state_manager is not None:
-                return self.service._state_manager.cache
-            else:
-                logger.error("Exposed service does not have a state manager.")
-                return {}
+            return self.state_manager.cache.cache
 
         # exposing custom.css file provided by user
         if self.css is not None: