from collections.abc import Mapping, Sequence
from typing import Any, Unpack
import pandas as pd
from ..offline.types import PlaceholderTranslations, SourcePlaceholderTranslations
from ..translator_typing import AbstractFetcherParams
from ..types import ID, IdType, SourceType
from ._abstract_fetcher import AbstractFetcher
from .types import FetchInstruction
[docs]
class MemoryFetcher(AbstractFetcher[SourceType, IdType]):
"""Fetch from memory.
Args:
data: A dict ``{source: PlaceholderTranslations}`` to fetch from.
return_all: If ``False``, return only the requested IDs and placeholders.
**kwargs: See :class:`.AbstractFetcher`.
"""
def __init__(
self,
data: (
SourcePlaceholderTranslations[SourceType]
| dict[SourceType, PlaceholderTranslations[SourceType]]
| dict[SourceType, pd.DataFrame]
| Mapping[SourceType, Mapping[str, Sequence[Any]]]
| None
) = None,
return_all: bool = True,
**kwargs: Unpack[AbstractFetcherParams[SourceType, IdType]],
) -> None:
super().__init__(**kwargs)
self._data: SourcePlaceholderTranslations[SourceType] = (
{} if not data else {source: PlaceholderTranslations.make(source, pht) for source, pht in data.items()}
)
self.return_all = return_all
def _initialize_sources(self, task_id: int) -> dict[SourceType, list[str]]: # noqa: ARG002
return {source: list(pht.placeholders) for source, pht in self._data.items()}
[docs]
def fetch_translations(self, instr: FetchInstruction[SourceType, IdType]) -> PlaceholderTranslations[SourceType]:
ret = self._data[instr.source]
if self.return_all:
return ret
placeholder_indices = [
ret.placeholders.index(placeholder)
for placeholder in instr.placeholders # requested
if placeholder in ret.placeholders # available
]
if instr.ids is None:
records = [[record[i] for i in placeholder_indices] for record in ret.records]
else:
ids = set(instr.ids)
records = [
[record[i] for i in placeholder_indices]
for record in ret.records
if record[ret.id_pos] in ids # crash on missing IDs
]
placeholders = tuple(ret.placeholders[i] for i in placeholder_indices)
id_pos = placeholders.index(ID) if ID in placeholders else -1
return PlaceholderTranslations(instr.source, placeholders, records, id_pos)
