yamily/yamily/yaml.py

import datetime  # pylint: disable=unused-import; doctest
import pathlib
import typing

import yaml

from yamily import Person


class _Loader(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: "_Loader", 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


class Dumper(yaml.SafeDumper):

    """
    >>> p = Person('alice')
    >>> p.name = 'Alice'
    >>> p.birth_date = datetime.date(1976, 2, 1)
    >>> print(yaml.dump(p, Dumper=Dumper))
    !person
    birth_date: 1976-02-01
    identifier: alice
    name: Alice
    <BLANKLINE>

    >>> p = Person('bob')
    >>> p.mother = Person('bob-mum')
    >>> p.father = Person('bob-father')
    >>> print(yaml.dump(p, Dumper=Dumper))
    !person
    father: !person
      identifier: bob-father
    identifier: bob
    mother: !person
      identifier: bob-mum
    <BLANKLINE>
    """

    # pylint: disable=too-many-ancestors

    def __init__(self, stream, **kwargs):
        super().__init__(stream, **kwargs)
        self.add_representer(Person, self._represent_person)

    @staticmethod
    def _represent_person(dumper: "_Dumper", person: Person) -> yaml.nodes.MappingNode:
        person_attrs = {"identifier": person.identifier}
        if person.name is not None:
            person_attrs["name"] = person.name
        if person.birth_date is not None:
            person_attrs["birth_date"] = person.birth_date
        if person.mother is not None:
            person_attrs["mother"] = person.mother
        if person.father is not None:
            person_attrs["father"] = person.father
        return dumper.represent_mapping("!person", person_attrs)


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=_Loader)