123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105 |
- import typing
- import graphviz
- from yamily import Person, PersonCollection
- 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 [shape=box]
- }
- {
- rank=same
- carol [shape=box]
- 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 [shape=box]
- }
- {
- rank=same
- alice [shape=box]
- }
- {
- rank=same
- 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")
- subgraph.node(person.identifier, shape="box") # , str(person))
- nodes.add(person)
- for coparent in collection.get_coparents(person):
- subgraph.node(coparent.identifier, shape="box")
- 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
|