Easier and more self documenting patterns for loading/saving PersistentHEADmaster

state via pickle and json.

1 files changed, 111 insertions, 5 deletions

diff --git a/persistent.py b/persistent.py
index 6cc444c..950471e 100644
--- a/persistent.py
+++ b/persistent.py
@@ -11,8 +11,11 @@ import datetime
 import enum
 import functools
 import logging
+import re
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Optional
+
+from overrides import overrides
 
 import file_utils
 
@@ -59,6 +62,113 @@ class Persistent(ABC):
         pass
 
 
+class FileBasedPersistent(Persistent):
+    """A Persistent that uses a file to save/load data and knows the conditions
+    under which the state should be saved/loaded."""
+
+    @staticmethod
+    @abstractmethod
+    def get_filename() -> str:
+        """Since this class saves/loads to/from a file, what's its full path?"""
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def should_we_save_data(filename: str) -> bool:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def should_we_load_data(filename: str) -> bool:
+        pass
+
+    @abstractmethod
+    def get_persistent_data(self) -> Any:
+        pass
+
+
+class PicklingFileBasedPersistent(FileBasedPersistent):
+    @classmethod
+    @overrides
+    def load(cls) -> Optional[Any]:
+        filename = cls.get_filename()
+        if cls.should_we_load_data(filename):
+            logger.debug('Attempting to load state from %s', filename)
+
+            import pickle
+
+            try:
+                with open(filename, 'rb') as rf:
+                    data = pickle.load(rf)
+                    return cls(data)
+
+            except Exception as e:
+                raise Exception(f'Failed to load {filename}.') from e
+        return None
+
+    @overrides
+    def save(self) -> bool:
+        filename = self.get_filename()
+        if self.should_we_save_data(filename):
+            logger.debug('Trying to save state in %s', filename)
+            try:
+                import pickle
+
+                with open(filename, 'wb') as wf:
+                    pickle.dump(self.get_persistent_data(), wf, pickle.HIGHEST_PROTOCOL)
+                return True
+            except Exception as e:
+                raise Exception(f'Failed to save to {filename}.') from e
+        return False
+
+
+class JsonFileBasedPersistent(FileBasedPersistent):
+    @classmethod
+    @overrides
+    def load(cls) -> Any:
+        filename = cls.get_filename()
+        if cls.should_we_load_data(filename):
+            logger.debug('Trying to load state from %s', filename)
+            import json
+
+            try:
+                with open(filename, 'r') as rf:
+                    lines = rf.readlines()
+
+                # This is probably bad... but I like comments
+                # in config files and JSON doesn't support them.  So
+                # pre-process the buffer to remove comments thus
+                # allowing people to add them.
+                buf = ''
+                for line in lines:
+                    line = re.sub(r'#.*$', '', line)
+                    buf += line
+
+                json_dict = json.loads(buf)
+                return cls(json_dict)
+
+            except Exception as e:
+                logger.exception(e)
+                raise Exception(f'Failed to load {filename}.') from e
+        return None
+
+    @overrides
+    def save(self) -> bool:
+        filename = self.get_filename()
+        if self.should_we_save_data(filename):
+            logger.debug('Trying to save state in %s', filename)
+            try:
+                import json
+
+                json_blob = json.dumps(self.get_persistent_data())
+                with open(filename, 'w') as wf:
+                    wf.writelines(json_blob)
+                return True
+            except Exception as e:
+                raise Exception(f'Failed to save to {filename}.') from e
+        return False
+
+
 def was_file_written_today(filename: str) -> bool:
     """Convenience wrapper around :meth:`was_file_written_within_n_seconds`.
 
@@ -225,10 +335,6 @@ class persistent_autoloaded_singleton(object):
         return _load
 
 
-# TODO: PicklingPersistant?
-# TODO: JsonConfigPersistant?
-
-
 if __name__ == '__main__':
     import doctest