"""Local caches for Atlas.""" from __future__ import annotations from collections import OrderedDict from typing import Iterable, Optional from app.models import AtlasEntity class EntityCache: def __init__(self, max_entries: int = 512): self.max_entries = max_entries self._data: OrderedDict[str, AtlasEntity] = OrderedDict() def _normalize_token(self, token: str | None) -> str: return str(token or "").strip().lower() def get(self, token: str | None) -> Optional[AtlasEntity]: key = self._normalize_token(token) if not key: return None entity = self._data.get(key) if entity is not None: self._data.move_to_end(key) return entity def store(self, entity: AtlasEntity, extra_tokens: Iterable[str] | None = None) -> None: tokens = set(extra_tokens or []) tokens.add(entity.canonical_label) tokens.add(entity.atlas_id) for alias in entity.aliases: tokens.add(alias.label) mid = entity.active_identifier("mid") if mid: tokens.add(mid) for token in tokens: key = self._normalize_token(token) if not key: continue self._data[key] = entity self._data.move_to_end(key) while len(self._data) > self.max_entries: self._data.popitem(last=False)