Source code for betty.project.extension.cotton_candy.search
"""
Provide Cotton Candy's search functionality.
"""
from __future__ import annotations
from abc import ABC
from asyncio import gather
from dataclasses import dataclass
from inspect import getmembers
from typing import TYPE_CHECKING, TypeVar, Generic, final
from typing_extensions import override
from betty.ancestry.file import File
from betty.ancestry.has_notes import HasNotes
from betty.ancestry.person import Person
from betty.ancestry.place import Place
from betty.ancestry.source import Source
from betty.locale.localizable import StaticTranslationsLocalizableAttr
from betty.model import Entity
from betty.privacy import is_private
from betty.typing import internal
if TYPE_CHECKING:
from betty.jinja2 import Environment
from betty.ancestry import Ancestry
from betty.locale.localizable import StaticTranslationsLocalizable
from betty.locale.localizer import Localizer
from betty.job import Context
from collections.abc import Iterable, Sequence
_EntityT = TypeVar("_EntityT", bound=Entity)
def _static_translations_to_text(
translations: StaticTranslationsLocalizable,
) -> set[str]:
return {
word
for translation in translations.translations.values()
for word in translation.strip().lower().split()
}
class _EntityTypeIndexer(Generic[_EntityT], ABC):
def text(self, entity: _EntityT) -> set[str]:
text = set()
# Each note is owner by a single other entity, so index it as part of that entity.
if isinstance(entity, HasNotes):
for note in entity.notes:
text.update(_static_translations_to_text(note.text))
for attr_name, class_attr_value in getmembers(type(entity)):
if isinstance(class_attr_value, StaticTranslationsLocalizableAttr):
text.update(_static_translations_to_text(getattr(entity, attr_name)))
return text
class _PersonIndexer(_EntityTypeIndexer[Person]):
@override
def text(self, entity: Person) -> set[str]:
text = super().text(entity)
for name in entity.names:
if name.individual is not None:
text.update(set(name.individual.lower().split()))
if name.affiliation is not None:
text.update(set(name.affiliation.lower().split()))
return text
class _PlaceIndexer(_EntityTypeIndexer[Place]):
@override
def text(self, entity: Place) -> set[str]:
text = super().text(entity)
for name in entity.names:
text.update(_static_translations_to_text(name.name))
return text
class _FileIndexer(_EntityTypeIndexer[File]):
pass
class _SourceIndexer(_EntityTypeIndexer[Source]):
pass
@final
@dataclass(frozen=True)
class _Entry:
text: set[str]
result: str
[docs]
@internal
class Index:
"""
Build search indexes.
"""
[docs]
async def build(self) -> Sequence[_Entry]:
"""
Build the search index.
"""
return [
entry
for entries in await gather(
self._build_entities(_PersonIndexer(), Person),
self._build_entities(_PlaceIndexer(), Place),
self._build_entities(_FileIndexer(), File),
self._build_entities(_SourceIndexer(), Source),
)
for entry in entries
if entry is not None
]