"""A share driver that stores data in memory """ from typing import Dict, List, Tuple from ..document import Document from ..exc import ReplicaIsClosedError from ..identity import Identity from ..path import Path from ..query import HistoryMode, Query from ..share import Share from . import Replica class InMemoryReplica(Replica): """In-memory Replica""" def __init__(self, share: Share, **driver_kwargs) -> None: self.share = share self._is_closed = False self._max_local_index = -1 # Local Index <=> Document pairs self._documents: List[Tuple[int, Document]] = {} @property def is_closed(self) -> bool: return self._is_closed def close(self, erase: bool = False) -> None: if self._is_closed: raise ReplicaIsClosedError() if erase: self._local_index = -1 self._documents = [] self._is_closed = True @property def max_local_index(self) -> int: if self._is_closed: raise ReplicaIsClosedError() return self._max_local_index def _get_all_docs(self) -> List[Document]: """Get all documents""" if self._is_closed: raise ReplicaIsClosedError() return [document for _, document in self._documents] def _get_latest_docs(self) -> List[Document]: """Get the latest version of each document""" if self._is_closed: raise ReplicaIsClosedError() docs_by_path: Dict[str, Document] = {} for document in self._documents: if ( str(document.path) not in docs_by_path or docs_by_path[str(document.path)].timestamp <= document.timestamp ): docs_by_path[str(document.path)] = document return list(docs_by_path.values()) def query_docs(self, query: Query) -> List[Document]: """Query a list of documents""" if self._is_closed: raise ReplicaIsClosedError() if query.history_mode == HistoryMode.ALL: docs = self._get_all_docs() else: docs = self._get_latest_docs() docs_to_local_index = { document: local_index for document, local_index in self.docs_by_local_index } return query({docs_to_local_index[document]: document for document in docs}) def upsert(self, new_document: Document) -> None: if self._is_closed: raise ReplicaIsClosedError() self._max_local_index += 1 self._documents = [ (local_index, document) for local_index, document in self._documents if document.author != new_document.author or document.path != new_document.path ] self._documents.append((self._local_index, new_document))