yamily/yamily/_graphviz.py

import typing

import graphviz

from yamily import Person, PersonCollection


def _add_person_node(graph: graphviz.dot.Digraph, person: Person) -> None:
    label = person.name or person.identifier
    graph.node(person.identifier, label=label, shape="box")


def digraph(collection: PersonCollection) -> graphviz.dot.Digraph:
    """
    >>> bob = Person('bob')
    >>> bob.father = Person('frank')

    >>> carol = Person('carol')
    >>> carol.mother = Person('grace')

    >>> alice = Person('alice')
    >>> alice.mother = carol
    >>> alice.father = bob

    >>> david = Person('david')
    >>> david.mother = carol
    >>> david.father = bob

    >>> collection = PersonCollection()
    >>> collection.add_person(alice)
    Person(alice)
    >>> collection.add_person(david)
    Person(david)
    >>> graph = digraph(collection)
    >>> print(graph.source)
    digraph yamily {
    	{
    		rank=same
    		grace [label=grace shape=box]
    	}
    	{
    		rank=same
    		carol [label=carol shape=box]
    		bob [label=bob shape=box]
    		"relation-bob-carol" [shape=point width=0]
    		carol -> "relation-bob-carol" [arrowhead=none constraint=False]
    		bob -> "relation-bob-carol" [arrowhead=none constraint=False]
    	}
    	{
    		rank=same
    		frank [label=frank shape=box]
    	}
    	{
    		rank=same
    		alice [label=alice shape=box]
    	}
    	{
    		rank=same
    		david [label=david shape=box]
    	}
    	grace -> carol
    	frank -> bob
    	"relation-bob-carol" -> alice
    	"relation-bob-carol" -> david
    }
    >>> graph.render('/tmp/yamily.gv')
    '/tmp/yamily.gv.pdf'
    """
    graph = graphviz.Digraph("yamily")
    nodes: typing.Set[Person] = set()
    parent_node_names: typing.Dict[typing.Tuple[Person, Person], str] = dict()
    for person in collection:
        if person in nodes:
            continue
        with graph.subgraph() as subgraph:
            subgraph.attr(rank="same")
            _add_person_node(subgraph, person)
            nodes.add(person)
            for coparent in collection.get_coparents(person):
                _add_person_node(subgraph, coparent)
                nodes.add(coparent)
                parents = tuple(sorted((person, coparent), key=lambda p: p.identifier))
                parent_node_name = "relation-{}-{}".format(
                    parents[0].identifier, parents[1].identifier
                )
                subgraph.node(parent_node_name, shape="point", width="0")
                subgraph.edge(
                    person.identifier,
                    parent_node_name,
                    constraint="False",
                    arrowhead="none",
                )
                subgraph.edge(
                    coparent.identifier,
                    parent_node_name,
                    constraint="False",
                    arrowhead="none",
                )
                parent_node_names[parents] = parent_node_name
                parent_node_names[tuple(parents[::-1])] = parent_node_name
    for child in collection:
        if child.mother is not None and child.father is not None:
            parents = sorted(child.parents, key=lambda p: p.identifier)
            parent_node_name = parent_node_names[tuple(parents)]
            graph.edge(parent_node_name, child.identifier)
        elif child.mother is not None:
            graph.edge(child.mother.identifier, child.identifier)
        elif child.father is not None:
            graph.edge(child.father.identifier, child.identifier)
    return graph