diff options
Diffstat (limited to 'persistent.py')
| -rw-r--r-- | persistent.py | 116 |
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 |
