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(
  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(60),
  86             persist_at_shutdown = False
  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.
 105         3. If the file doesn't exist or the predicate indicates that
 106            the persisted state cannot be reused, MyComplexObject's
 107            __init__ will be invoked and will be expected to fully
 108            initialize the instance.
 109         4. At program exit time, depending on the value of the
 110            persist_at_shutdown parameter, the state of MyComplexObject
 111            will be written to disk using the same filename so that
 112            future instances may potentially reuse saved state.  Note
 113            that the state that is persisted is the state at program
 114            exit time.
 115
 116     """
 117     def __init__(
 118             self,
 119             filename: str,
 120             *,
 121             may_reuse_persisted: Optional[Callable[[datetime.datetime], bool]] = None,
 122             persist_at_shutdown: PersistAtShutdown = PersistAtShutdown.NEVER):
 123         self.filename = filename
 124         self.may_reuse_persisted = may_reuse_persisted
 125         self.persist_at_shutdown = persist_at_shutdown
 126         self.instance = None
 127
 128     def __call__(self, cls):
 129         @functools.wraps(cls)
 130         def _load(*args, **kwargs):
 131
 132             # If class has already been loaded, act like a singleton
 133             # and return a reference to the one and only instance in
 134             # memory.
 135             if self.instance is not None:
 136                 logger.debug(
 137                     f'Returning already instantiated singleton instance of {cls.__name__}.'
 138                 )
 139                 return self.instance
 140
 141             was_loaded_from_disk = False
 142             if file_utils.does_file_exist(self.filename):
 143                 cache_mtime_dt = file_utils.get_file_mtime_as_datetime(
 144                     self.filename
 145                 )
 146                 now = datetime.datetime.now()
 147                 if (
 148                         self.may_reuse_persisted is not None and
 149                         self.may_reuse_persisted(cache_mtime_dt)
 150                 ):
 151                     logger.debug(
 152                         f'Attempting to load from persisted cache (mtime={cache_mtime_dt}, {now-cache_mtime_dt} ago)')
 153                     if not self.load():
 154                         logger.warning('Loading from cache failed?!')
 155                         assert self.instance is None
 156                     else:
 157                         assert self.instance is not None
 158                         was_loaded_from_disk = True
 159
 160             if self.instance is None:
 161                 logger.debug(
 162                     f'Attempting to instantiate {cls.__name__} directly.'
 163                 )
 164                 self.instance = cls(*args, **kwargs)
 165                 was_loaded_from_disk = False
 166
 167             assert self.instance is not None
 168             if (
 169                     self.persist_at_shutdown is PersistAtShutdown.ALWAYS or
 170                     (
 171                         not was_loaded_from_disk and
 172                         self.persist_at_shutdown is PersistAtShutdown.IF_NOT_INITIALIZED_FROM_DISK
 173                     )
 174             ):
 175                 atexit.register(self.save)
 176             return self.instance
 177         return _load
 178
 179     def load(self) -> bool:
 180         try:
 181             with open(self.filename, 'rb') as f:
 182                 self.instance = dill.load(f)
 183                 return True
 184         except Exception:
 185             self.instance = None
 186             return False
 187         return False
 188
 189     def save(self) -> bool:
 190         if self.instance is not None:
 191             logger.debug(
 192                 f'Attempting to save {type(self.instance).__name__} to file {self.filename}'
 193             )
 194             try:
 195                 with open(self.filename, 'wb') as f:
 196                     dill.dump(self.instance, f, dill.HIGHEST_PROTOCOL)
 197                 return True
 198             except Exception:
 199                 return False
 200         return False