Extract functionality into an internal search API

betty/project/extension/_theme/search.py

@@ -0,0 +1,233 @@
+"""
+Provide Cotton Candy's search functionality.
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC
+from asyncio import gather
+from dataclasses import dataclass
+from inspect import getmembers
+from typing import TYPE_CHECKING, TypeVar, Generic, final, cast
+
+import aiofiles
+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 StaticTranslationsLocalizable, Localizable
+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.machine_name import MachineName
+    from betty.project import Project
+    from betty.jinja2 import Environment
+    from betty.ancestry import Ancestry
+    from betty.locale.localizer import Localizer
+    from betty.job import Context
+    from collections.abc import Iterable, Sequence
+
+_EntityT = TypeVar("_EntityT", bound=Entity)
+
+
+async def generate_search_index(
+    project: Project,
+    result_container_template: Localizable,
+    results_container_template: Localizable,
+    *,
+    job_context: Context,
+) -> None:
+    await gather(
+        *(
+            _generate_search_index_for_locale(
+                project,
+                result_container_template,
+                results_container_template,
+                locale,
+                job_context=job_context,
+            )
+            for locale in project.configuration.locales
+        )
+    )
+
+
+async def _generate_search_index_for_locale(
+    project: Project,
+    result_container_template: Localizable,
+    results_container_template: Localizable,
+    locale: str,
+    *,
+    job_context: Context,
+) -> None:
+    localizers = await project.localizers
+    localizer = await localizers.get(locale)
+    search_index = {
+        "resultContainerTemplate": result_container_template.localize(localizer),
+        "resultsContainerTemplate": results_container_template.localize(localizer),
+        "index": [
+            {
+                "entityTypeId": entry.entity_type_id,
+                "text": " ".join(entry.text),
+                "result": entry.result,
+            }
+            for entry in await Index(
+                project.ancestry,
+                await project.jinja2_environment,
+                job_context,
+                localizer,
+            ).build()
+        ],
+    }
+    search_index_json = json.dumps(search_index)
+    async with aiofiles.open(
+        project.configuration.localize_www_directory_path(locale) / "search-index.json",
+        mode="w",
+    ) as f:
+        await f.write(search_index_json)
+
+
+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, localizer: Localizer, 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(
+                        cast(StaticTranslationsLocalizable, getattr(entity, attr_name))
+                    )
+                )
+
+        return text
+
+
+class _PersonIndexer(_EntityTypeIndexer[Person]):
+    @override
+    def text(self, localizer: Localizer, entity: Person) -> set[str]:
+        text = super().text(localizer, 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, localizer: Localizer, entity: Place) -> set[str]:
+        text = super().text(localizer, entity)
+        for name in entity.names:
+            text.update(_static_translations_to_text(name.name))
+        return text
+
+
+class _FileIndexer(_EntityTypeIndexer[File]):
+    @override
+    def text(self, localizer: Localizer, entity: File) -> set[str]:
+        text = super().text(localizer, entity)
+        text.update(entity.label.localize(localizer).strip().lower().split())
+        return text
+
+
+class _SourceIndexer(_EntityTypeIndexer[Source]):
+    pass
+
+
+@final
+@dataclass(frozen=True)
+class _Entry:
+    entity_type_id: MachineName
+    result: str
+    text: set[str]
+
+
+@internal
+class Index:
+    """
+    Build search indexes.
+    """
+
+    def __init__(
+        self,
+        ancestry: Ancestry,
+        jinja2_environment: Environment,
+        job_context: Context | None,
+        localizer: Localizer,
+    ):
+        self._ancestry = ancestry
+        self._jinja2_environment = jinja2_environment
+        self._job_context = job_context
+        self._localizer = localizer
+
+    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
+        ]
+
+    async def _build_entities(
+        self, indexer: _EntityTypeIndexer[_EntityT], entity_type: type[_EntityT]
+    ) -> Iterable[_Entry | None]:
+        return await gather(
+            *(
+                self._build_entity(indexer, entity)
+                for entity in self._ancestry[entity_type]
+            )
+        )
+
+    async def _build_entity(
+        self, indexer: _EntityTypeIndexer[_EntityT], entity: _EntityT
+    ) -> _Entry | None:
+        if is_private(entity):
+            return None
+        text = indexer.text(self._localizer, entity)
+        if not text:
+            return None
+        return _Entry(entity.plugin_id(), await self._render_entity(entity), text)
+
+    async def _render_entity(self, entity: Entity) -> str:
+        return await self._jinja2_environment.select_template(
+            [
+                f"search/result--{entity.plugin_id()}.html.j2",
+                "search/result.html.j2",
+            ]
+        ).render_async(
+            {
+                "job_context": self._job_context,
+                "localizer": self._localizer,
+                "entity": entity,
+            }
+        )

betty/project/extension/cotton_candy/__init__.py

@@ -4,26 +4,20 @@
 
 from __future__ import annotations
 
-import json
-from asyncio import gather
 from pathlib import Path
 from typing import TYPE_CHECKING, final, Self
 
-import aiofiles
 from typing_extensions import override
 
 from betty.html import CssProvider
-from betty.jinja2 import (
-    Jinja2Provider,
-    Filters,
-)
-from betty.locale.localizable import _, static
+from betty.jinja2 import Jinja2Provider, Filters
+from betty.locale.localizable import _, static, plain
 from betty.os import link_or_copy
 from betty.plugin import ShorthandPluginBase
 from betty.project.extension import ConfigurableExtension, Theme, Extension
 from betty.project.extension._theme import jinja2_filters
+from betty.project.extension._theme.search import generate_search_index
 from betty.project.extension.cotton_candy.config import CottonCandyConfiguration
-from betty.project.extension.cotton_candy.search import Index
 from betty.project.extension.maps import Maps
 from betty.project.extension.trees import Trees
 from betty.project.extension.webpack import Webpack
@@ -37,17 +31,17 @@
     from betty.event_dispatcher import EventHandlerRegistry
     from collections.abc import Sequence
 
-_RESULT_CONTAINER_TEMPLATE = """
+_RESULT_CONTAINER_TEMPLATE = plain("""
 <li class="search-result">
     {{{ betty-search-result }}}
 </li>
-"""
+""")
 
-_RESULTS_CONTAINER_TEMPLATE = """
+_RESULTS_CONTAINER_TEMPLATE = plain("""
 <ul id="search-results" class="nav-secondary">
     {{{ betty-search-results }}}
 </ul>
-"""
+""")
 
 
 async def _generate_logo(event: GenerateSiteEvent) -> None:
@@ -57,42 +51,14 @@ async def _generate_logo(event: GenerateSiteEvent) -> None:
 
 
 async def _generate_search_index(event: GenerateSiteEvent) -> None:
-    await gather(
-        *(
-            _generate_search_index_for_locale(event, locale)
-            for locale in event.project.configuration.locales
-        )
+    await generate_search_index(
+        event.project,
+        _RESULT_CONTAINER_TEMPLATE,
+        _RESULTS_CONTAINER_TEMPLATE,
+        job_context=event.job_context,
     )
 
 
-async def _generate_search_index_for_locale(
-    event: GenerateSiteEvent, locale: str
-) -> None:
-    project = event.project
-    localizers = await project.localizers
-    localizer = await localizers.get(locale)
-    search_index = {
-        "resultContainerTemplate": _RESULT_CONTAINER_TEMPLATE,
-        "resultsContainerTemplate": _RESULTS_CONTAINER_TEMPLATE,
-        "index": [
-            {"text": " ".join(entry.text), "result": entry.result}
-            for entry in await Index(
-                project.ancestry,
-                await project.jinja2_environment,
-                event.job_context,
-                localizer,
-            ).build()
-        ],
-    }
-    search_index_json = json.dumps(search_index)
-    async with aiofiles.open(
-        event.project.configuration.localize_www_directory_path(locale)
-        / "search-index.json",
-        mode="w",
-    ) as f:
-        await f.write(search_index_json)
-
-
 @final
 class CottonCandy(
     ShorthandPluginBase,