persistent.py

   1 #!/usr/bin/env python3
   2
   3 from abc import ABC, abstractmethod
   4 import atexit
   5 import datetime
   6 import enum
   7 import functools
   8 import logging
   9 from typing import Callable, Optional
  10
  11 import dill
  12
  13 import file_utils
  14
  15 logger = logging.getLogger(__name__)
  16
  17
  18 class Persistent(ABC):
  19     """
  20     A base class of an object with a load/save method.
  21     """
  22     @abstractmethod
  23     def save(self):
  24         pass
  25
  26     @abstractmethod
  27     def load(self):
  28         pass
  29
  30
  31 def reuse_if_mtime_is_today() -> Callable[[datetime.datetime], bool]:
  32     """
  33     A helper that returns a lambda appropriate for use in the
  34     persistent_autoload_singleton decorator's may_reuse_persisted
  35     parameter that allows persisted state to be reused as long as it
  36     was persisted on the same day as the load.
  37
  38     """
  39     now = datetime.datetime.now()
  40     return lambda dt: (
  41         dt.month == now.month and
  42         dt.day == now.day and
  43         dt.year == now.year
  44     )
  45
  46
  47 def reuse_if_mtime_less_than_limit_sec(
  48         limit_seconds: int
  49 ) -> Callable[[datetime.datetime], bool]:
  50     """
  51     A helper that returns a lambda appropriate for use in the
  52     persistent_autoload_singleton decorator's may_reuse_persisted
  53     parameter that allows persisted state to be reused as long as it
  54     was persisted within the past limit_seconds.
  55
  56     """
  57     now = datetime.datetime.now()
  58     return lambda dt: (now - dt).total_seconds() <= limit_seconds
  59
  60
  61 class PersistAtShutdown(enum.Enum):
  62     """
  63     An enum to describe the conditions under which state is persisted
  64     to disk.  See details below.
  65
  66     """
  67     NEVER = 0,
  68     IF_NOT_INITIALIZED_FROM_DISK = 1,
  69     ALWAYS = 2,
  70
  71
  72 class persistent_autoload_singleton(Persistent):
  73     """This class is meant to be used as a decorator around a class that:
  74
  75         1. Is a singleton; one global instance per python program.
  76         2. Has a complex state that is initialized fully by __init__()
  77         3. Would benefit from caching said state on disk and reloading
  78            it on future invokations rather than recomputing and
  79            reinitializing.
  80
  81     Here's and example usage pattern:
  82
  83         @persistent_autoload_singleton(
  84             filename = "my_cache_file.bin",
  85             may_reuse_persisted = reuse_if_mtime_less_than_limit_sec(60),
  86             persist_at_shutdown = PersistAtShutdown.IF_NOT_INITIALIZED_FROM_DISK,
  87         )
  88         class MyComplexObject(object):
  89             def __init__(self, ...):
  90                 # do a bunch of work to fully initialize this instance
  91
  92             def another_method(self, ...):
  93                 # use the state stored in this instance to do some work
  94
  95     What does this do, exactly?
  96
  97         1. Anytime you attempt to instantiate MyComplexObject you will
  98            get the same instance.  This class is now a singleton.
  99         2. The first time you attempt to instantiate MyComplexObject
 100            the wrapper scaffolding will check "my_cache_file.bin".  If
 101            it exists and any may_reuse_persisted predicate indicates
 102            that reusing persisted state is allowed, we will skip the
 103            call to __init__ and return an unpickled instance read from
 104            the disk file.  In the example above the predicate allows
 105            reuse of saved state if it is <= 60s old.
 106         3. If the file doesn't exist or the predicate indicates that
 107            the persisted state cannot be reused (e.g. too stale),
 108            MyComplexObject's __init__ will be invoked and will be
 109            expected to fully initialize the instance.
 110         4. At program exit time, depending on the value of the
 111            persist_at_shutdown parameter, the state of MyComplexObject
 112            will be written to disk using the same filename so that
 113            future instances may potentially reuse saved state.  Note
 114            that the state that is persisted is the state at program
 115            exit time.  In the example above this parameter indicates
 116            that we should persist state so long as we were not
 117            initialized from cached state on disk.
 118
 119     """
 120     def __init__(
 121             self,
 122             filename: str,
 123             *,
 124             may_reuse_persisted: Optional[Callable[[datetime.datetime], bool]] = None,
 125             persist_at_shutdown: PersistAtShutdown = PersistAtShutdown.NEVER):
 126         self.filename = filename
 127         self.may_reuse_persisted = may_reuse_persisted
 128         self.persist_at_shutdown = persist_at_shutdown
 129         self.instance = None
 130
 131     def __call__(self, cls):
 132         @functools.wraps(cls)
 133         def _load(*args, **kwargs):
 134
 135             # If class has already been loaded, act like a singleton
 136             # and return a reference to the one and only instance in
 137             # memory.
 138             if self.instance is not None:
 139                 logger.debug(
 140                     f'Returning already instantiated singleton instance of {cls.__name__}.'
 141                 )
 142                 return self.instance
 143
 144             was_loaded_from_disk = False
 145             if file_utils.does_file_exist(self.filename):
 146                 cache_mtime_dt = file_utils.get_file_mtime_as_datetime(
 147                     self.filename
 148                 )
 149                 now = datetime.datetime.now()
 150                 if (
 151                         self.may_reuse_persisted is not None and
 152                         self.may_reuse_persisted(cache_mtime_dt)
 153                 ):
 154                     logger.debug(
 155                         f'Attempting to load from persisted cache (mtime={cache_mtime_dt}, {now-cache_mtime_dt} ago)')
 156                     if not self.load():
 157                         logger.warning('Loading from cache failed?!')
 158                         assert self.instance is None
 159                     else:
 160                         assert self.instance is not None
 161                         was_loaded_from_disk = True
 162
 163             if self.instance is None:
 164                 logger.debug(
 165                     f'Attempting to instantiate {cls.__name__} directly.'
 166                 )
 167                 self.instance = cls(*args, **kwargs)
 168                 was_loaded_from_disk = False
 169
 170             assert self.instance is not None
 171             if (
 172                     self.persist_at_shutdown is PersistAtShutdown.ALWAYS or
 173                     (
 174                         not was_loaded_from_disk and
 175                         self.persist_at_shutdown is PersistAtShutdown.IF_NOT_INITIALIZED_FROM_DISK
 176                     )
 177             ):
 178                 atexit.register(self.save)
 179             return self.instance
 180         return _load
 181
 182     def load(self) -> bool:
 183         try:
 184             with open(self.filename, 'rb') as f:
 185                 self.instance = dill.load(f)
 186                 return True
 187         except Exception:
 188             self.instance = None
 189             return False
 190         return False
 191
 192     def save(self) -> bool:
 193         if self.instance is not None:
 194             logger.debug(
 195                 f'Attempting to save {type(self.instance).__name__} to file {self.filename}'
 196             )
 197             try:
 198                 with open(self.filename, 'wb') as f:
 199                     dill.dump(self.instance, f, dill.HIGHEST_PROTOCOL)
 200                 return True
 201             except Exception:
 202                 return False
 203         return False