From a9b98b675dedddd116866bbc79583dcd83faa523 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mose=20M=C3=BCller?= <mosemueller@gmail.com>
Date: Wed, 2 Aug 2023 12:06:21 +0200
Subject: [PATCH] feat: adding support for saving the status of a DataService
 class

By passing a the "filename" keyword to the DataService init function,
one can save the state of the instance in this file. Quitting the service
(Ctrl+C) will dump the state of the instance in the file with given
filename. When starting the instance, it will load these values again.
---
 .../data_service/data_service.py              | 62 ++++++++++++++++++-
 src/pyDataInterface/server/server.py          |  3 +
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/src/pyDataInterface/data_service/data_service.py b/src/pyDataInterface/data_service/data_service.py
index 1735dc5..441e7c7 100644
--- a/src/pyDataInterface/data_service/data_service.py
+++ b/src/pyDataInterface/data_service/data_service.py
@@ -180,18 +180,76 @@ class DataService(rpyc.Service, TaskManager):
     the "subject" and the callback functions as the "observers".
     """
 
-    def __init__(self) -> None:
+    def __init__(self, filename: Optional[str] = None) -> None:
         TaskManager.__init__(self)
-        DataServiceSerializer.__init__(self, "serialized.json")
         self.__root__: "DataService" = self
         """Keep track of the root object. This helps to filter the emission of
         notifications. This overwrite the TaksManager's __root__ attribute."""
 
         self._callbacks: set[Callable[[str, Any], None]] = set()
+        self._filename: Optional[str] = filename
 
         self._register_callbacks()
         self.__check_instance_classes()
         self._initialised = True
+        self._load_values_from_json()
+
+    def _load_values_from_json(self) -> None:
+        if self._filename is not None:
+            # Check if the file specified by the filename exists
+            if os.path.exists(self._filename):
+                with open(self._filename, "r") as f:
+                    # Load JSON data from file and update class attributes with these
+                    # values
+                    self.set_attributes_from_serialized_representation(
+                        cast(dict[str, Any], json.load(f))
+                    )
+
+    def write_to_file(self) -> None:
+        """
+        Serialize the DataService instance and write it to a JSON file.
+
+        Args:
+            filename (str): The name of the file to write to.
+        """
+        if self._filename is not None:
+            with open(self._filename, "w") as f:
+                json.dump(self.serialize(), f, indent=4)
+        else:
+            logger.error(
+                f"Class {self.__class__.__name__} was not initialised with a filename. "
+                'Skipping "write_to_file"...'
+            )
+
+    def set_attributes_from_serialized_representation(
+        self, serialized_representation: dict[str, Any]
+    ) -> None:
+        # Traverse the serialized representation and set the attributes of the class
+        for path, value in generate_paths_and_values_from_serialized_DataService(
+            serialized_representation
+        ).items():
+            # Split the path into elements
+            parent_path, attr_name = f"DataService.{path}".rsplit(".", 1)
+
+            if isinstance(value, list):
+                for index, item in enumerate(value):
+                    update_DataService_by_path(
+                        self,
+                        {
+                            "name": f"{attr_name}[{index}]",
+                            "parent_path": parent_path,
+                            "value": item,
+                        },
+                    )
+            else:
+                update_DataService_by_path(
+                    self,
+                    {
+                        "name": attr_name,
+                        "parent_path": parent_path,
+                        "value": value,
+                    },
+                )
 
     def __setattr__(self, __name: str, __value: Any) -> None:
         current_value = getattr(self, __name, None)
diff --git a/src/pyDataInterface/server/server.py b/src/pyDataInterface/server/server.py
index a1778fc..054fc1c 100644
--- a/src/pyDataInterface/server/server.py
+++ b/src/pyDataInterface/server/server.py
@@ -173,6 +173,9 @@ class Server:
     async def shutdown(self) -> None:
         logger.info("Shutting down")
 
+        logger.info(f"Saving data to {self._service._filename}.")
+        self._service.write_to_file()
+
         await self.__cancel_servers()
         await self.__cancel_tasks()