src/pyutils/persistent.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2023, Scott Gasch
   4
   5 """
   6 This module defines a class hierarchy (base class :class:`Persistent`) and
   7 a decorator (`@persistent_autoloaded_singleton`) that can be used to create
   8 objects that load and save their state from some external storage location
   9 automatically, optionally and conditionally.
  10
  11 A :class:`Persistent` is just a class with a :meth:`Persistent.load` and
  12 :meth:`Persistent.save` method.   Various subclasses such as
  13 :class:`JsonFileBasedPersistent` and :class:`PicklingFileBasedPersistent`
  14 define these methods to, save data in a particular format.  The details
  15 of where and whether to save are left to your code to decide by implementing
  16 interface methods like :meth:`FileBasedPersistent.get_filename` and
  17 :meth:`FileBasedPersistent.should_we_load_data`.
  18
  19 This module inculdes some helpers to make deciding whether to load persisted
  20 state easier such as :meth:`was_file_written_today` and
  21 :meth:`was_file_written_within_n_seconds`.
  22
  23 :class:`Persistent` classes are good for things backed by persisted
  24 state that is loaded all or most of the time.  For example, the high
  25 score list of a game, the configuration settings of a tool,
  26 etc... Really anything that wants to save/load state from storage and
  27 not bother with the plumbing to do so.
  28 """
  29
  30 import atexit
  31 import datetime
  32 import enum
  33 import functools
  34 import logging
  35 import re
  36 from abc import ABC, abstractmethod
  37 from typing import Any, Optional
  38
  39 from overrides import overrides
  40
  41 from pyutils.files import file_utils
  42
  43 logger = logging.getLogger(__name__)
  44
  45
  46 class Persistent(ABC):
  47     """
  48     A base class of an object with a load/save method.  Classes that are
  49     decorated with :code:`@persistent_autoloaded_singleton` should subclass
  50     this and implement their :meth:`save` and :meth:`load` methods.
  51     """
  52
  53     @abstractmethod
  54     def save(self) -> bool:
  55         """
  56         Save this thing somewhere that you'll remember when someone calls
  57         :meth:`load` later on in a way that makes sense to your code.
  58         """
  59         pass
  60
  61     @classmethod
  62     @abstractmethod
  63     def load(cls) -> Any:
  64         """Load this thing from somewhere and give back an instance which
  65         will become the global singleton and which may (see
  66         below) be saved (via :meth:`save`) at program exit time.
  67
  68         Oh, in case this is handy, here's a reminder how to write a
  69         factory method that doesn't call the c'tor in python::
  70
  71             @classmethod
  72             def load_from_somewhere(cls, somewhere):
  73                 # Note: __new__ does not call __init__.
  74                 obj = cls.__new__(cls)
  75
  76                 # Don't forget to call any polymorphic base class initializers
  77                 super(MyClass, obj).__init__()
  78
  79                 # Load the piece(s) of obj that you want to from somewhere.
  80                 obj._state = load_from_somewhere(somewhere)
  81                 return obj
  82         """
  83         pass
  84
  85
  86 class FileBasedPersistent(Persistent):
  87     """A :class:`Persistent` subclass that uses a file to save/load
  88     data and knows the conditions under which the state should be
  89     saved/loaded.
  90     """
  91
  92     @staticmethod
  93     @abstractmethod
  94     def get_filename() -> str:
  95         """
  96         Returns:
  97             The full path of the file in which we are saving/loading data.
  98         """
  99         pass
 100
 101     @staticmethod
 102     @abstractmethod
 103     def should_we_save_data(filename: str) -> bool:
 104         """
 105         Returns:
 106             True if we should save our state now or False otherwise.
 107         """
 108         pass
 109
 110     @staticmethod
 111     @abstractmethod
 112     def should_we_load_data(filename: str) -> bool:
 113         """
 114         Returns:
 115             True if we should load persisted state now or False otherwise.
 116         """
 117         pass
 118
 119     @abstractmethod
 120     def get_persistent_data(self) -> Any:
 121         """
 122         Returns:
 123             The raw state data read from the filesystem.  Can be any format.
 124         """
 125         pass
 126
 127
 128 class PicklingFileBasedPersistent(FileBasedPersistent):
 129     """
 130     A class that stores its state in a file as pickled Python objects.
 131
 132     Example usage::
 133
 134         import persistent
 135
 136         @persistent.persistent_autoloaded_singleton()
 137         class MyClass(persistent.PicklingFileBasedPersistent):
 138             def __init__(self, data: Optional[Whatever]):
 139                 if data:
 140                     # initialize state from data
 141                 else:
 142                     # if desired, initialize an "empty" object with new state.
 143
 144             @staticmethod
 145             @overrides
 146             def get_filename() -> str:
 147                 return "/path/to/where/you/want/to/save/data.bin"
 148
 149             @staticmethod
 150             @overrides
 151             def should_we_save_data(filename: str) -> bool:
 152                 return true_if_we_should_save_the_data_this_time()
 153
 154             @staticmethod
 155             @overrides
 156             def should_we_load_data(filename: str) -> bool:
 157                 return persistent.was_file_written_within_n_seconds(whatever)
 158
 159         # Persistent will handle the plumbing to instantiate your class from its
 160         # persisted state iff the :meth:`should_we_load_data` says it's ok to.  It
 161         # will also persist the current in-memory state to disk at program exit iff
 162         # the :meth:`should_we_save_data` methods says to.
 163         c = MyClass()
 164
 165     """
 166
 167     @abstractmethod
 168     def __init__(self, data: Optional[Any] = None):
 169         """You should override this."""
 170         pass
 171
 172     @classmethod
 173     @overrides
 174     def load(cls) -> Optional[Any]:
 175         filename = cls.get_filename()
 176         if cls.should_we_load_data(filename):
 177             logger.debug("Attempting to load state from %s", filename)
 178             assert file_utils.is_readable(filename)
 179
 180             import pickle
 181
 182             try:
 183                 with open(filename, "rb") as rf:
 184                     data = pickle.load(rf)
 185                     return cls(data)
 186
 187             except Exception as e:
 188                 raise Exception(f"Failed to load {filename}.") from e
 189         return None
 190
 191     @overrides
 192     def save(self) -> bool:
 193         filename = self.get_filename()
 194         if self.should_we_save_data(filename):
 195             logger.debug("Trying to save state in %s", filename)
 196             try:
 197                 import pickle
 198
 199                 with file_utils.CreateFileWithMode(filename, 0o600, "wb") as wf:
 200                     pickle.dump(self.get_persistent_data(), wf, pickle.HIGHEST_PROTOCOL)
 201                 return True
 202             except Exception as e:
 203                 raise Exception(f"Failed to save to {filename}.") from e
 204         return False
 205
 206
 207 class JsonFileBasedPersistent(FileBasedPersistent):
 208     """A class that stores its state in a JSON format file.
 209
 210     Example usage::
 211
 212         import persistent
 213
 214         @persistent.persistent_autoloaded_singleton()
 215         class MyClass(persistent.JsonFileBasedPersistent):
 216             def __init__(self, data: Optional[dict[str, Any]]):
 217                 # load already deserialized the JSON data for you; it's
 218                 # a "cooked" JSON dict of string -> values, lists, dicts,
 219                 # etc...
 220                 if data:
 221                     #initialize youself from data...
 222                 else:
 223                     # if desired, initialize an empty state object
 224                     # when json_data isn't provided.
 225
 226             @staticmethod
 227             @overrides
 228             def get_filename() -> str:
 229                 return "/path/to/where/you/want/to/save/data.json"
 230
 231             @staticmethod
 232             @overrides
 233             def should_we_save_data(filename: str) -> bool:
 234                 return true_if_we_should_save_the_data_this_time()
 235
 236             @staticmethod
 237             @overrides
 238             def should_we_load_data(filename: str) -> bool:
 239                 return persistent.was_file_written_within_n_seconds(whatever)
 240
 241         # Persistent will handle the plumbing to instantiate your
 242         # class from its persisted state iff the
 243         # :meth:`should_we_load_data` says it's ok to.  It will also
 244         # persist the current in memory state to disk at program exit
 245         # iff the :meth:`should_we_save_data methods` says to.
 246         c = MyClass()
 247     """
 248
 249     @abstractmethod
 250     def __init__(self, data: Optional[Any]):
 251         """You should override this."""
 252         pass
 253
 254     @classmethod
 255     @overrides
 256     def load(cls) -> Any:
 257         filename = cls.get_filename()
 258         if cls.should_we_load_data(filename):
 259             logger.debug("Trying to load state from %s", filename)
 260             import json
 261
 262             try:
 263                 with open(filename, "r") as rf:
 264                     lines = rf.readlines()
 265
 266                 # This is probably bad... but I like comments
 267                 # in config files and JSON doesn't support them.  So
 268                 # pre-process the buffer to remove comments thus
 269                 # allowing people to add them.
 270                 buf = ""
 271                 for line in lines:
 272                     line = re.sub(r"#.*$", "", line)
 273                     buf += line
 274                 json_dict = json.loads(buf)
 275                 return cls(json_dict)
 276
 277             except Exception as e:
 278                 logger.exception(
 279                     "Failed to load path %s; raising an exception", filename
 280                 )
 281                 raise Exception(f"Failed to load {filename}.") from e
 282         return None
 283
 284     @overrides
 285     def save(self) -> bool:
 286         filename = self.get_filename()
 287         if self.should_we_save_data(filename):
 288             logger.debug("Trying to save state in %s", filename)
 289             try:
 290                 import json
 291
 292                 json_blob = json.dumps(self.get_persistent_data())
 293                 with open(filename, "w") as wf:
 294                     wf.writelines(json_blob)
 295                 return True
 296             except Exception as e:
 297                 raise Exception(f"Failed to save to {filename}.") from e
 298         return False
 299
 300
 301 def was_file_written_today(filename: str) -> bool:
 302     """Convenience wrapper around :meth:`was_file_written_within_n_seconds`.
 303
 304     Args:
 305         filename: path / filename to check
 306
 307     Returns:
 308         True if filename was written today.
 309
 310     >>> import os
 311     >>> filename = f'/tmp/testing_persistent_py_{os.getpid()}'
 312     >>> os.system(f'touch {filename}')
 313     0
 314     >>> was_file_written_today(filename)
 315     True
 316     >>> os.system(f'touch -d 1974-04-15T01:02:03.99 {filename}')
 317     0
 318     >>> was_file_written_today(filename)
 319     False
 320     >>> os.system(f'/bin/rm -f {filename}')
 321     0
 322     >>> was_file_written_today(filename)
 323     False
 324     """
 325     if not file_utils.does_file_exist(filename):
 326         return False
 327
 328     mtime = file_utils.get_file_mtime_as_datetime(filename)
 329     assert mtime is not None
 330     now = datetime.datetime.now()
 331     return mtime.month == now.month and mtime.day == now.day and mtime.year == now.year
 332
 333
 334 def was_file_written_within_n_seconds(
 335     filename: str,
 336     limit_seconds: int,
 337 ) -> bool:
 338     """Helper for determining persisted state staleness.
 339
 340     Args:
 341         filename: the filename to check
 342         limit_seconds: how fresh, in seconds, it must be
 343
 344     Returns:
 345         True if filename was written within the past limit_seconds
 346         or False otherwise (or on error).
 347
 348     >>> import os
 349     >>> filename = f'/tmp/testing_persistent_py_{os.getpid()}'
 350     >>> os.system(f'touch {filename}')
 351     0
 352     >>> was_file_written_within_n_seconds(filename, 60)
 353     True
 354     >>> import time
 355     >>> time.sleep(2.0)
 356     >>> was_file_written_within_n_seconds(filename, 2)
 357     False
 358     >>> os.system(f'/bin/rm -f {filename}')
 359     0
 360     >>> was_file_written_within_n_seconds(filename, 60)
 361     False
 362     """
 363
 364     if not file_utils.does_file_exist(filename):
 365         return False
 366
 367     mtime = file_utils.get_file_mtime_as_datetime(filename)
 368     assert mtime is not None
 369     now = datetime.datetime.now()
 370     return (now - mtime).total_seconds() <= limit_seconds
 371
 372
 373 class PersistAtShutdown(enum.Enum):
 374     """
 375     An enum to describe the conditions under which state is persisted
 376     to disk.  This is passed as an argument to the decorator below and
 377     is used to indicate when to call :meth:`save` on a :class:`Persistent`
 378     subclass.
 379
 380     * NEVER: never call :meth:`save`
 381     * IF_NOT_LOADED: call :meth:`save` as long as we did not successfully
 382       :meth:`load` its state.
 383     * ALWAYS: always call :meth:`save`
 384     """
 385
 386     NEVER = (0,)
 387     IF_NOT_LOADED = (1,)
 388     ALWAYS = (2,)
 389
 390
 391 class persistent_autoloaded_singleton(object):
 392     """A decorator that can be applied to a :class:`Persistent` subclass
 393     (i.e.  a class with :meth:`save` and :meth:`load` methods.  The
 394     decorator will intercept attempts to instantiate the class via
 395     it's c'tor and, instead, invoke the class' :meth:`load` to give it a
 396     chance to read state from somewhere persistent (disk, db,
 397     whatever).  Subsequent calls to construt instances of the wrapped
 398     class will return a single, global instance (i.e. the wrapped
 399     class is must be a singleton).
 400
 401     If :meth:`load` fails (returns None), the c'tor is invoked with the
 402     original args as a fallback.
 403
 404     Based upon the value of the optional argument
 405     :code:`persist_at_shutdown` argument, (NEVER, IF_NOT_LOADED,
 406     ALWAYS), the :meth:`save` method of the class will be invoked just
 407     before program shutdown to give the class a chance to save its
 408     state somewhere.
 409
 410     .. note::
 411         The implementations of :meth:`save` and :meth:`load` and where the
 412         class persists its state are details left to the :class:`Persistent`
 413         implementation.  Essentially this decorator just handles the
 414         plumbing of calling your save/load and appropriate times and
 415         creates a transparent global singleton whose state can be
 416         persisted between runs.
 417
 418     """
 419
 420     def __init__(
 421         self,
 422         *,
 423         persist_at_shutdown: PersistAtShutdown = PersistAtShutdown.IF_NOT_LOADED,
 424     ):
 425         self.persist_at_shutdown = persist_at_shutdown
 426         self.instance = None
 427
 428     def __call__(self, cls: Persistent):
 429         @functools.wraps(cls)  # type: ignore
 430         def _load(*args, **kwargs):
 431
 432             # If class has already been loaded, act like a singleton
 433             # and return a reference to the one and only instance in
 434             # memory.
 435             if self.instance is not None:
 436                 logger.debug(
 437                     "Returning already instantiated singleton instance of %s.",
 438                     cls.__name__,
 439                 )
 440                 return self.instance
 441
 442             # Otherwise, try to load it from persisted state.
 443             was_loaded = False
 444             logger.debug("Attempting to load %s from persisted state.", cls.__name__)
 445             self.instance = cls.load()
 446             if not self.instance:
 447                 msg = "Loading from cache failed."
 448                 logger.warning(msg)
 449                 logger.debug("Attempting to instantiate %s directly.", cls.__name__)
 450                 self.instance = cls(*args, **kwargs)
 451             else:
 452                 logger.debug(
 453                     "Class %s was loaded from persisted state successfully.",
 454                     cls.__name__,
 455                 )
 456                 was_loaded = True
 457
 458             assert self.instance is not None
 459
 460             if self.persist_at_shutdown is PersistAtShutdown.ALWAYS or (
 461                 not was_loaded
 462                 and self.persist_at_shutdown is PersistAtShutdown.IF_NOT_LOADED
 463             ):
 464                 logger.debug(
 465                     "Scheduling a deferred called to save at process shutdown time."
 466                 )
 467                 atexit.register(self.instance.save)
 468             return self.instance
 469
 470         return _load
 471
 472
 473 if __name__ == "__main__":
 474     import doctest
 475
 476     doctest.testmod()