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