yamily/family_tree_yaml/yaml.py

import pathlib
import typing

import yaml
from family_tree_yaml import Person


class _YamlLoader(yaml.SafeLoader):

    # pylint: disable=too-many-ancestors

    def __init__(self, stream):
        super().__init__(stream)
        self.add_constructor("!person", self._construct_person)

    @staticmethod
    def _construct_person(
        loader: "_YamlLoader", node: yaml.nodes.MappingNode
    ) -> Person:
        (person_attrs,) = loader.construct_yaml_map(node)
        person = Person(person_attrs["identifier"])
        if "name" in person_attrs:
            person.name = person_attrs["name"]
        if "birth_date" in person_attrs:
            person.birth_date = person_attrs["birth_date"]
        if "mother" in person_attrs:
            person.mother = Person(person_attrs["mother"])
        if "father" in person_attrs:
            person.father = Person(person_attrs["father"])
        return person


def read(yaml_path: typing.Union[str, pathlib.Path]) -> Person:
    """
    >>> read('persons/erika-mustermann.yml')
    Person(erika-mustermann, Erika Mustermann, *1957-08-12)
    >>> maxl = read('persons/max-mustermann.yml')
    >>> maxl
    Person(max-mustermann, Max Mustermann, *1976-02-01)
    >>> maxl.mother
    Person(erika-mustermann)
    >>> maxl.father
    Person(thomas-mustermann)
    """
    if isinstance(yaml_path, str):
        return read(pathlib.Path(yaml_path))
    with yaml_path.open("r") as yaml_file:
        return yaml.load(yaml_file, Loader=_YamlLoader)