yamily/yamily/_graphviz.py

import graphviz

from yamily import Person, PersonCollection  # pylint: disable=unused-import; doctest


def digraph(collection: PersonCollection) -> graphviz.dot.Digraph:
    """
    >>> alice = Person('alice')
    >>> alice.name = 'Alice'
    >>> alice.mother = Person('carol')
    >>> alice.father = Person('bob')
    >>> collection = PersonCollection()
    >>> collection.add_person(alice)
    Person(alice, Alice)
    >>> graph = digraph(collection)
    >>> print(graph.source)
    digraph yamily {
    	carol [label=unnamed]
    	bob [label=unnamed]
    	alice [label=Alice]
    	carol -> alice
    	bob -> alice
    }
    >>> graph.render('/tmp/yamily.gv')
    '/tmp/yamily.gv.pdf'
    """
    graph = graphviz.Digraph("yamily")
    for person in collection:
        graph.node(person.identifier, str(person))
    for child in collection:
        for parent in [child.mother, child.father]:
            if parent is not None:
                graph.edge(parent.identifier, child.identifier)
    return graph