yamily/yamily/_cli.py

import argparse
import pathlib
import sys
import typing

import yaml

import yamily
import yamily._graphviz
import yamily.yaml


def _read(path: pathlib.Path) -> typing.Iterator[yamily.Person]:
    if path.is_dir():
        for file_path in path.glob("**/*.yml"):
            for person in _read(file_path):
                yield person
    else:
        with path.open("r") as yaml_file:
            yield yaml.load(yaml_file, Loader=yamily.yaml.Loader)


def _list() -> None:
    argparser = argparse.ArgumentParser(
        description="Recursively find yamily family tree members and print as YAML list."
    )
    argparser.add_argument(
        "paths", nargs="+", metavar="path", help="path to yamily .yml file or folder"
    )
    args = argparser.parse_args()
    collection = yamily.PersonCollection()
    for path in args.paths:
        for person in _read(pathlib.Path(path)):
            collection.add_person(person)
    yaml.dump(
        sorted(collection, key=lambda p: p.identifier),
        sys.stdout,
        Dumper=yamily.yaml.Dumper,
        default_flow_style=False,
        allow_unicode=True,
    )


def _dot() -> None:
    argparser = argparse.ArgumentParser(
        description="Create family tree in DOT (graphviz) format. "
        "Recursively looks for *.yml files containing family members in yamily format."
    )
    argparser.add_argument(
        "paths", nargs="+", metavar="path", help="path to yamily .yml file or folder"
    )
    args = argparser.parse_args()
    collection = yamily.PersonCollection()
    for path in args.paths:
        for person in _read(pathlib.Path(path)):
            collection.add_person(person)
    graph = yamily._graphviz.digraph(collection)  # pylint: disable=protected-access
    print(graph.source)