Source code for stackholm.storages.optimized_list.optimized_list_state

from contextlib import suppress
from typing import (
    Dict,
    List,
    Optional,
)

from stackholm.context import Context
from stackholm.state import State


__all__ = (
    'OptimizedListState',
)


[docs]class OptimizedListState(State): """ List-based state implementation. """ context_sequence: int """ Sequence number of the contexts. """ contexts: List[Context] """ List of contexts. """ checkpoint_sequences: Dict[str, int] """ Mapping of checkpoint keys to their sequence numbers. """ checkpoint_indexes: Dict[str, List[int]] """ Mapping of checkpoint keys to their context indexes. """ checkpoint_optimization_mapping: Dict[str, Dict[int, int]] """ Mapping of checkpoint keys to another mapping of context indexes to their checkpoint indexes. """
[docs] def __init__(self) -> None: self.context_sequence = -1 self.contexts = [] self.checkpoint_sequences = {} self.checkpoint_indexes = {} self.checkpoint_optimization_mapping = {}
[docs] def push_context( self, context: Context, ) -> int: """ Appends a context to the list. Returns the index of the context appended. :param context: The context to be appended. """ self.context_sequence += 1 self.contexts.append(context) return self.context_sequence
[docs] def pop_context( self, index: int = -1, ) -> Optional[Context]: """ Pops and returns a context from the list. If the list is empty, returns `None`. :param index: The index of the context to be popped. """ self.context_sequence -= 1 with suppress(IndexError): return self.contexts.pop(index) return None
[docs] def get_last_context(self) -> Optional[Context]: """ Returns the last context in the list. If the list is empty, returns `None`. """ with suppress(IndexError): return self.contexts[-1] return None
[docs] def add_checkpoint( self, key: str, context_index: int, ) -> None: """ 1. Creates a new checkpoint sequence number if it does not exist. 2. Creates a partition if it does not exist. 3. Creates a mapping if it does not exist, for supporting delete operations with O(1) time complexity. 4. Adds the context index to the partition. :param key: The partition key. :param context_index: The index of the context. """ if key not in self.checkpoint_sequences: self.checkpoint_sequences[key] = -1 if key not in self.checkpoint_optimization_mapping: self.checkpoint_optimization_mapping[key] = {} if key not in self.checkpoint_indexes: self.checkpoint_indexes[key] = [] self.checkpoint_sequences[key] += 1 checkpoint_index = self.checkpoint_sequences[key] self.checkpoint_optimization_mapping[key][context_index] = checkpoint_index self.checkpoint_indexes[key].append(context_index)
[docs] def remove_checkpoint( self, key: str, context_index: int, ) -> None: """ 1. If partition does not exist, returns. 2. Finds the checkpoint index using the optimization mapping, this has O(1). 3. It should be impossible with the high-level usage but if somehow the checkpoint index is not found, `list.index()` is used to find the checkpoint index. 4. Removes the context index from the partition. :param key: The partition key. :param context_index: The index of the context. """ checkpoint_index: Optional[int] = None if key in self.checkpoint_optimization_mapping: checkpoint_index = self.checkpoint_optimization_mapping[key].pop(context_index, None) if not self.checkpoint_optimization_mapping[key]: self.checkpoint_optimization_mapping.pop(key, None) if key in self.checkpoint_indexes and checkpoint_index is None: with suppress(ValueError): checkpoint_index = self.checkpoint_indexes[key].index(context_index) if checkpoint_index is not None: if key in self.checkpoint_sequences: self.checkpoint_sequences[key] -= 1 if self.checkpoint_sequences[key] < 0: self.checkpoint_sequences.pop(key, None) self.checkpoint_indexes[key].pop(checkpoint_index) if not self.checkpoint_indexes[key]: self.checkpoint_indexes.pop(key, None)
[docs] def get_nearest_checkpoint( self, key: str, ) -> Optional[Context]: """ Returns the latest context in the partition with the given partition key. If the partition does not exist, returns `None`. :param key: The partition key. """ with suppress(KeyError, IndexError): return self.contexts[self.checkpoint_indexes[key][-1]] return None