src/pyutils/persistent.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """A :class:`Persistent` is just a class with a load and save method.  This
   6 module defines the :class:`Persistent` base and a decorator that can be used to
   7 create a persistent singleton that autoloads and autosaves."""
   8
   9 import atexit
  10 import datetime
  11 import enum
  12 import functools
  13 import logging
  14 import re
  15 from abc import ABC, abstractmethod
  16 from typing import Any, Optional
  17
  18 from overrides import overrides
  19
  20 from pyutils.files import file_utils
  21
  22 logger = logging.getLogger(__name__)
  23
  24
  25 class Persistent(ABC):
  26     """
  27     A base class of an object with a load/save method.  Classes that are
  28     decorated with :code:`@persistent_autoloaded_singleton` should subclass
  29     this and implement their :meth:`save` and :meth:`load` methods.
  30     """
  31
  32     @abstractmethod
  33     def save(self) -> bool:
  34         """
  35         Save this thing somewhere that you'll remember when someone calls
  36         :meth:`load` later on in a way that makes sense to your code.
  37         """
  38         pass
  39
  40     @classmethod
  41     @abstractmethod
  42     def load(cls) -> Any:
  43         """Load this thing from somewhere and give back an instance which
  44         will become the global singleton and which may (see
  45         below) be saved (via :meth:`save`) at program exit time.
  46
  47         Oh, in case this is handy, here's a reminder how to write a
  48         factory method that doesn't call the c'tor in python::
  49
  50             @classmethod
  51             def load_from_somewhere(cls, somewhere):
  52                 # Note: __new__ does not call __init__.
  53                 obj = cls.__new__(cls)
  54
  55                 # Don't forget to call any polymorphic base class initializers
  56                 super(MyClass, obj).__init__()
  57
  58                 # Load the piece(s) of obj that you want to from somewhere.
  59                 obj._state = load_from_somewhere(somewhere)
  60                 return obj
  61         """
  62         pass
  63
  64
  65 class FileBasedPersistent(Persistent):
  66     """A Persistent that uses a file to save/load data and knows the conditions
  67     under which the state should be saved/loaded."""
  68
  69     @staticmethod
  70     @abstractmethod
  71     def get_filename() -> str:
  72         """Since this class saves/loads to/from a file, what's its full path?"""
  73         pass
  74
  75     @staticmethod
  76     @abstractmethod
  77     def should_we_save_data(filename: str) -> bool:
  78         pass
  79
  80     @staticmethod
  81     @abstractmethod
  82     def should_we_load_data(filename: str) -> bool:
  83         pass
  84
  85     @abstractmethod
  86     def get_persistent_data(self) -> Any:
  87         pass
  88
  89
  90 class PicklingFileBasedPersistent(FileBasedPersistent):
  91     @classmethod
  92     @overrides
  93     def load(cls) -> Optional[Any]:
  94         filename = cls.get_filename()
  95         if cls.should_we_load_data(filename):
  96             logger.debug('Attempting to load state from %s', filename)
  97             assert file_utils.file_is_readable(filename)
  98
  99             import pickle
 100
 101             try:
 102                 with open(filename, 'rb') as rf:
 103                     data = pickle.load(rf)
 104                     return cls(data)
 105
 106             except Exception as e:
 107                 raise Exception(f'Failed to load {filename}.') from e
 108         return None
 109
 110     @overrides
 111     def save(self) -> bool:
 112         filename = self.get_filename()
 113         if self.should_we_save_data(filename):
 114             logger.debug('Trying to save state in %s', filename)
 115             try:
 116                 import pickle
 117
 118                 with open(filename, 'wb') as wf:
 119                     pickle.dump(self.get_persistent_data(), wf, pickle.HIGHEST_PROTOCOL)
 120                 return True
 121             except Exception as e:
 122                 raise Exception(f'Failed to save to {filename}.') from e
 123         return False
 124
 125
 126 class JsonFileBasedPersistent(FileBasedPersistent):
 127     @classmethod
 128     @overrides
 129     def load(cls) -> Any:
 130         filename = cls.get_filename()
 131         if cls.should_we_load_data(filename):
 132             logger.debug('Trying to load state from %s', filename)
 133             import json
 134
 135             try:
 136                 with open(filename, 'r') as rf:
 137                     lines = rf.readlines()
 138
 139                 # This is probably bad... but I like comments
 140                 # in config files and JSON doesn't support them.  So
 141                 # pre-process the buffer to remove comments thus
 142                 # allowing people to add them.
 143                 buf = ''
 144                 for line in lines:
 145                     line = re.sub(r'#.*$', '', line)
 146                     buf += line
 147
 148                 json_dict = json.loads(buf)
 149                 return cls(json_dict)
 150
 151             except Exception as e:
 152                 logger.exception(e)
 153                 raise Exception(f'Failed to load {filename}.') from e
 154         return None
 155
 156     @overrides
 157     def save(self) -> bool:
 158         filename = self.get_filename()
 159         if self.should_we_save_data(filename):
 160             logger.debug('Trying to save state in %s', filename)
 161             try:
 162                 import json
 163
 164                 json_blob = json.dumps(self.get_persistent_data())
 165                 with open(filename, 'w') as wf:
 166                     wf.writelines(json_blob)
 167                 return True
 168             except Exception as e:
 169                 raise Exception(f'Failed to save to {filename}.') from e
 170         return False
 171
 172
 173 def was_file_written_today(filename: str) -> bool:
 174     """Convenience wrapper around :meth:`was_file_written_within_n_seconds`.
 175
 176     Args:
 177         filename: filename to check
 178
 179     Returns:
 180         True if filename was written today.
 181
 182     >>> import os
 183     >>> filename = f'/tmp/testing_persistent_py_{os.getpid()}'
 184     >>> os.system(f'touch {filename}')
 185     0
 186     >>> was_file_written_today(filename)
 187     True
 188     >>> os.system(f'touch -d 1974-04-15T01:02:03.99 {filename}')
 189     0
 190     >>> was_file_written_today(filename)
 191     False
 192     >>> os.system(f'/bin/rm -f {filename}')
 193     0
 194     >>> was_file_written_today(filename)
 195     False
 196     """
 197
 198     if not file_utils.does_file_exist(filename):
 199         return False
 200
 201     mtime = file_utils.get_file_mtime_as_datetime(filename)
 202     assert mtime is not None
 203     now = datetime.datetime.now()
 204     return mtime.month == now.month and mtime.day == now.day and mtime.year == now.year
 205
 206
 207 def was_file_written_within_n_seconds(
 208     filename: str,
 209     limit_seconds: int,
 210 ) -> bool:
 211     """Helper for determining persisted state staleness.
 212
 213     Args:
 214         filename: the filename to check
 215         limit_seconds: how fresh, in seconds, it must be
 216
 217     Returns:
 218         True if filename was written within the past limit_seconds
 219         or False otherwise (or on error).
 220
 221     >>> import os
 222     >>> filename = f'/tmp/testing_persistent_py_{os.getpid()}'
 223     >>> os.system(f'touch {filename}')
 224     0
 225     >>> was_file_written_within_n_seconds(filename, 60)
 226     True
 227     >>> import time
 228     >>> time.sleep(2.0)
 229     >>> was_file_written_within_n_seconds(filename, 2)
 230     False
 231     >>> os.system(f'/bin/rm -f {filename}')
 232     0
 233     >>> was_file_written_within_n_seconds(filename, 60)
 234     False
 235     """
 236
 237     if not file_utils.does_file_exist(filename):
 238         return False
 239
 240     mtime = file_utils.get_file_mtime_as_datetime(filename)
 241     assert mtime is not None
 242     now = datetime.datetime.now()
 243     return (now - mtime).total_seconds() <= limit_seconds
 244
 245
 246 class PersistAtShutdown(enum.Enum):
 247     """
 248     An enum to describe the conditions under which state is persisted
 249     to disk.  This is passed as an argument to the decorator below and
 250     is used to indicate when to call :meth:`save` on a :class:`Persistent`
 251     subclass.
 252
 253     * NEVER: never call :meth:`save`
 254     * IF_NOT_LOADED: call :meth:`save` as long as we did not successfully
 255       :meth:`load` its state.
 256     * ALWAYS: always call :meth:`save`
 257     """
 258
 259     NEVER = (0,)
 260     IF_NOT_LOADED = (1,)
 261     ALWAYS = (2,)
 262
 263
 264 class persistent_autoloaded_singleton(object):
 265     """A decorator that can be applied to a :class:`Persistent` subclass
 266     (i.e.  a class with :meth:`save` and :meth:`load` methods.  The
 267     decorator will intercept attempts to instantiate the class via
 268     it's c'tor and, instead, invoke the class' :meth:`load` to give it a
 269     chance to read state from somewhere persistent (disk, db,
 270     whatever).  Subsequent calls to construt instances of the wrapped
 271     class will return a single, global instance (i.e. the wrapped
 272     class is a singleton).
 273
 274     If :meth:`load` fails (returns None), the c'tor is invoked with the
 275     original args as a fallback.
 276
 277     Based upon the value of the optional argument
 278     :code:`persist_at_shutdown` argument, (NEVER, IF_NOT_LOADED,
 279     ALWAYS), the :meth:`save` method of the class will be invoked just
 280     before program shutdown to give the class a chance to save its
 281     state somewhere.
 282
 283     .. note::
 284         The implementations of :meth:`save` and :meth:`load` and where the
 285         class persists its state are details left to the :class:`Persistent`
 286         implementation.  Essentially this decorator just handles the
 287         plumbing of calling your save/load and appropriate times and
 288         creates a transparent global singleton whose state can be
 289         persisted between runs.
 290
 291     """
 292
 293     def __init__(
 294         self,
 295         *,
 296         persist_at_shutdown: PersistAtShutdown = PersistAtShutdown.IF_NOT_LOADED,
 297     ):
 298         self.persist_at_shutdown = persist_at_shutdown
 299         self.instance = None
 300
 301     def __call__(self, cls: Persistent):
 302         @functools.wraps(cls)  # type: ignore
 303         def _load(*args, **kwargs):
 304
 305             # If class has already been loaded, act like a singleton
 306             # and return a reference to the one and only instance in
 307             # memory.
 308             if self.instance is not None:
 309                 logger.debug(
 310                     'Returning already instantiated singleton instance of %s.',
 311                     cls.__name__,
 312                 )
 313                 return self.instance
 314
 315             # Otherwise, try to load it from persisted state.
 316             was_loaded = False
 317             logger.debug('Attempting to load %s from persisted state.', cls.__name__)
 318             self.instance = cls.load()
 319             if not self.instance:
 320                 msg = 'Loading from cache failed.'
 321                 logger.warning(msg)
 322                 logger.debug('Attempting to instantiate %s directly.', cls.__name__)
 323                 self.instance = cls(*args, **kwargs)
 324             else:
 325                 logger.debug(
 326                     'Class %s was loaded from persisted state successfully.',
 327                     cls.__name__,
 328                 )
 329                 was_loaded = True
 330
 331             assert self.instance is not None
 332
 333             if self.persist_at_shutdown is PersistAtShutdown.ALWAYS or (
 334                 not was_loaded
 335                 and self.persist_at_shutdown is PersistAtShutdown.IF_NOT_LOADED
 336             ):
 337                 logger.debug(
 338                     'Scheduling a deferred called to save at process shutdown time.'
 339                 )
 340                 atexit.register(self.instance.save)
 341             return self.instance
 342
 343         return _load
 344
 345
 346 if __name__ == '__main__':
 347     import doctest
 348
 349     doctest.testmod()