betty.ancestry - Betty documentation

Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for betty.ancestry

    """
    Provide Betty's main data model.
    """

    from __future__ import annotations

    from contextlib import contextmanager
    from typing import Iterable, final, TYPE_CHECKING

    from betty.model import Entity
    from betty.model.association import AssociationRegistry
    from betty.model.collections import MultipleTypesEntityCollection

    if TYPE_CHECKING:
        from betty.plugin import PluginIdToTypeMapping
        from collections.abc import Iterator

    @final
    class Ancestry(MultipleTypesEntityCollection[Entity]):
        """
        An ancestry contains all the entities of a single family tree/genealogical data set.
        """

        def __init__(
            self,
            *entities: Entity,
            entity_type_id_to_type_mapping: PluginIdToTypeMapping[Entity],
        ):
            super().__init__(
                *entities, entity_type_id_to_type_mapping=entity_type_id_to_type_mapping
            )
            self._check_graph = True

        @contextmanager
        def unchecked(self) -> Iterator[None]:
            """
            Disable the addition entities' associates when adding those entities to the ancestry.

            It is the caller's responsibility to ensure all associates are added to the ancestry.
            If this is done, using this context manager improves performance.
            """
            self._check_graph = False
            try:
                yield
            finally:
                self._check_graph = True

        def _on_add(self, *entities: Entity) -> None:
            super()._on_add(*entities)
            if self._check_graph:
                self.add(*self._get_associates(*entities))

        def _get_associates(self, *entities: Entity) -> Iterable[Entity]:
            for entity in entities:
                for association in AssociationRegistry.get_all_associations(entity):
                    for associate in association.get_associates(entity):
                        yield associate