Source code for stackholm.context

from contextlib import ContextDecorator
from types import TracebackType
from typing import (
    Any,
    Dict,
    Optional,
    TYPE_CHECKING,
    Type,
    TypeVar,
)

from stackholm.exceptions import (
    ContextIsNotActive,
    NoContextIsActive,
)


if TYPE_CHECKING:
    from stackholm.storage import Storage


__all__ = (
    'Context',
)


VALUE_T = TypeVar('VALUE_T', bound=object)


[docs]class Context(ContextDecorator): """ Context is the base class for all context classes created via the :meth:`.Storage.create_context_class` method. .. caution:: Do not create instances of this class directly, otherwise they will not be bound to a :class:`.Storage` instance. """ _storage: 'Storage' """ :class:`.Storage` instance the class is bound to. """ _index: Optional[int] """ Index of the context in the :attr:`.Storage.contexts` list. """ _data: Dict[str, Any] """ Dictionary of context data. """
[docs] @classmethod def get_current(cls) -> Optional['Context']: """ Returns the latest context in the stack. If the stack is empty returns `None`. """ return cls._storage.get_last_context()
[docs] @classmethod def get_nearest_checkpoint( cls, key: str, ) -> Optional['Context']: """ Returns the latest context that contains the given key from the partitioned stack. If the partitioned stack is empty, no context contains the key in this case, returns `None`. :param key: The partition key here is the key used to store a value in a context. """ return cls._storage.get_nearest_checkpoint(key)
[docs] @classmethod def get_checkpoint_value( cls, key: str, default: Optional[VALUE_T] = None, ) -> Optional[VALUE_T]: """ Returns the value of the given key in the nearest checkpoint context. :param key: The key of the value to be returned. :param default: The default value to be returned in case the key is not found. (Default: `None`) """ context = cls.get_nearest_checkpoint(key) if context is None: return default return context.data.get(key, default)
[docs] @classmethod def set_checkpoint_value( cls, key: str, value: Any, ) -> None: """ Sets the value of the given key in the current context. If no context is active, raises :class:`.NoContextIsActive`. .. admonition:: Note After setting a checkpoint value in a context, the context becomes a checkpoint for the specified key. :param key: The key of the value to be set. :param value: The value to be set. """ context = cls.get_current() if context is None: raise NoContextIsActive() context.data[key] = value cls._storage.add_checkpoint(key, context.index)
[docs] @classmethod def pop_checkpoint_value( cls, key: str, default: Optional[VALUE_T] = None, ) -> Optional[VALUE_T]: """ Returns the value of the given key in the nearest checkpoint context. :param key: The key of the value to be returned. :param default: The default value to be returned in case the key is not found. (Default: `None`) """ context = cls.get_nearest_checkpoint(key) if context is None: return default cls._storage.remove_checkpoint(key, context.index) return context.data.pop(key, default)
[docs] @classmethod def reset_checkpoint_value( cls, key: str, ) -> None: """ Deletes the values of the given key by clearing all the checkpoint markers in the partitioned stack. .. caution:: The time complexity of this operation is O(n), where n is the number of checkpoints in the partitioned stack. :param key: The key of the values to be deleted. """ while cls._storage.get_nearest_checkpoint(key) is not None: cls.pop_checkpoint_value(key)
[docs] def __init__(self) -> None: self._index = None self._data = {}
def __enter__(self) -> 'Context': return self.activate() def __exit__( self, exception_type: Optional[Type[BaseException]], exception: Optional[BaseException], traceback: Optional[TracebackType], ) -> None: self.deactivate() @property def storage(self) -> 'Storage': """ Returns the :class:`.Storage` instance used by this context's class. """ storage = getattr(self.__class__, '_storage', None) assert storage is not None, 'Context class must be bound to a storage.' # noqa return storage @property def index(self) -> int: """ Returns the index of this context in the stack. .. warning:: Calling this property before the context is activated raises :class:`.ContextIsNotActive`. """ if self._index is None: raise ContextIsNotActive(self) return self._index @property def is_active(self) -> bool: """ Returns `True` if this context is active, `False` otherwise. """ return self._index is not None @property def data(self) -> Dict[str, Any]: """ Returns the data of this context. """ return self._data
[docs] def get_value( self, key: str, default: VALUE_T = None, ) -> VALUE_T: """ Returns the value of the given key from this context. :param key: The key of the value to be returned. :param default: The default value to be returned in case the key is not found. (Default: `None`) """ return self.data.get(key, default)
[docs] def set_value( self, key: str, value: Any, ) -> None: """ Sets the value of the given key in this context. :param key: The key of the value to be set. :param value: The value to be set. """ self._data[key] = value
[docs] def activate(self) -> 'Context': """ Activates this context by pushing it to the stack. .. admonition:: Note This method is called automatically when entering the context. .. code-block:: python with context: # context is activated here. ... """ if not self.is_active: self._index = self.storage.push_context(self) return self
[docs] def deactivate(self) -> None: """ Deactivates this context by popping it from the stack. .. caution:: After deactivating a context, its index is invalidated by being set to `None`. Moreover, the sequence number of contexts in the storage decreases for reusing the index. Unlike a DBMS, this behavior is needed for list-based algorithms and performance reasons. Since all the write operations must be performed in synchronous way, we need to optimize the performance for read/delete operations only. .. admonition:: Note This method is called automatically when exiting the context. .. code-block:: python with context: ... # context is deactivated here. """ if self.is_active: for key in self.data.keys(): self.storage.remove_checkpoint(key, self.index) self.storage.pop_context(self.index) self._index = None