path: root/persistent.py

#!/usr/bin/env python3

from abc import ABC, abstractmethod
import atexit
import datetime
import enum
import functools
import logging
from typing import Callable, Optional

import dill

import file_utils

logger = logging.getLogger(__name__)


class Persistent(ABC):
    """
    A base class of an object with a load/save method.
    """
    @abstractmethod
    def save(self):
        pass

    @abstractmethod
    def load(self):
        pass


def reuse_if_mtime_is_today() -> Callable[[datetime.datetime], bool]:
    """
    A helper that returns a lambda appropriate for use in the
    persistent_autoloaded_singleton decorator's may_reuse_persisted
    parameter that allows persisted state to be reused as long as it
    was persisted on the same day as the load.

    """
    now = datetime.datetime.now()
    return lambda dt: (
        dt.month == now.month and
        dt.day == now.day and
        dt.year == now.year
    )


def reuse_if_mtime_less_than_limit_sec(
        limit_seconds: int
) -> Callable[[datetime.datetime], bool]:
    """
    A helper that returns a lambda appropriate for use in the
    persistent_autoloaded_singleton decorator's may_reuse_persisted
    parameter that allows persisted state to be reused as long as it
    was persisted within the past limit_seconds.

    """
    now = datetime.datetime.now()
    return lambda dt: (now - dt).total_seconds() <= limit_seconds


def dont_reuse_persisted_state_force_refresh(
) -> Callable[[datetime.datetime], bool]:
    return lambda dt: False


class PersistAtShutdown(enum.Enum):
    """
    An enum to describe the conditions under which state is persisted
    to disk.  See details below.

    """
    NEVER = 0,
    IF_NOT_INITIALIZED_FROM_DISK = 1,
    ALWAYS = 2,


class persistent_autoloaded_singleton(Persistent):
    """This class is meant to be used as a decorator around a class that:

        1. Is a singleton; one global instance per python program.
        2. Has a complex state that is initialized fully by __init__()
        3. Would benefit from caching said state on disk and reloading
           it on future invokations rather than recomputing and
           reinitializing.

    Here's and example usage pattern:

        @persistent_autoloaded_singleton(
            filename = "my_cache_file.bin",
            may_reuse_persisted = reuse_if_mtime_less_than_limit_sec(60),
            persist_at_shutdown = PersistAtShutdown.IF_NOT_INITIALIZED_FROM_DISK,
        )
        class MyComplexObject(object):
            def __init__(self, ...):
                # do a bunch of work to fully initialize this instance

            def another_method(self, ...):
                # use the state stored in this instance to do some work

    What does this do, exactly?

        1. Anytime you attempt to instantiate MyComplexObject you will
           get the same instance.  This class is now a singleton.
        2. The first time you attempt to instantiate MyComplexObject
           the wrapper scaffolding will check "my_cache_file.bin".  If
           it exists and any may_reuse_persisted predicate indicates
           that reusing persisted state is allowed, we will skip the
           call to __init__ and return an unpickled instance read from
           the disk file.  In the example above the predicate allows
           reuse of saved state if it is <= 60s old.
        3. If the file doesn't exist or the predicate indicates that
           the persisted state cannot be reused (e.g. too stale),
           MyComplexObject's __init__ will be invoked and will be
           expected to fully initialize the instance.
        4. At program exit time, depending on the value of the
           persist_at_shutdown parameter, the state of MyComplexObject
           will be written to disk using the same filename so that
           future instances may potentially reuse saved state.  Note
           that the state that is persisted is the state at program
           exit time.  In the example above this parameter indicates
           that we should persist state so long as we were not
           initialized from cached state on disk.

    """
    def __init__(
            self,
            filename: str,
            *,
            may_reuse_persisted: Optional[Callable[[datetime.datetime], bool]] = None,
            persist_at_shutdown: PersistAtShutdown = PersistAtShutdown.NEVER):
        self.filename = filename
        self.may_reuse_persisted = may_reuse_persisted
        self.persist_at_shutdown = persist_at_shutdown
        self.instance = None

    def __call__(self, cls):
        @functools.wraps(cls)
        def _load(*args, **kwargs):

            # If class has already been loaded, act like a singleton
            # and return a reference to the one and only instance in
            # memory.
            if self.instance is not None:
                logger.debug(
                    f'Returning already instantiated singleton instance of {cls.__name__}.'
                )
                return self.instance

            was_loaded_from_disk = False
            if file_utils.does_file_exist(self.filename):
                cache_mtime_dt = file_utils.get_file_mtime_as_datetime(
                    self.filename
                )
                now = datetime.datetime.now()
                if (
                        self.may_reuse_persisted is not None and
                        self.may_reuse_persisted(cache_mtime_dt)
                ):
                    logger.debug(
                        f'Attempting to load from persisted cache (mtime={cache_mtime_dt}, {now-cache_mtime_dt} ago)')
                    if not self.load():
                        logger.warning('Loading from cache failed?!')
                        assert self.instance is None
                    else:
                        assert self.instance is not None
                        was_loaded_from_disk = True

            if self.instance is None:
                logger.debug(
                    f'Attempting to instantiate {cls.__name__} directly.'
                )
                self.instance = cls(*args, **kwargs)
                was_loaded_from_disk = False

            assert self.instance is not None
            if (
                    self.persist_at_shutdown is PersistAtShutdown.ALWAYS or
                    (
                        not was_loaded_from_disk and
                        self.persist_at_shutdown is PersistAtShutdown.IF_NOT_INITIALIZED_FROM_DISK
                    )
            ):
                atexit.register(self.save)
            return self.instance
        return _load

    def load(self) -> bool:
        try:
            with open(self.filename, 'rb') as f:
                self.instance = dill.load(f)
                return True
        except Exception:
            self.instance = None
            return False
        return False

    def save(self) -> bool:
        if self.instance is not None:
            logger.debug(
                f'Attempting to save {type(self.instance).__name__} to file {self.filename}'
            )
            try:
                with open(self.filename, 'wb') as f:
                    dill.dump(self.instance, f, dill.HIGHEST_PROTOCOL)
                return True
            except Exception:
                return False
        return False