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