Browse Source

graphviz: parents on same rank

Fabian Peter Hammerle 4 years ago
parent
commit
6aa634c7d9
2 changed files with 96 additions and 27 deletions
  1. 11 13
      yamily/__init__.py
  2. 85 14
      yamily/_graphviz.py

+ 11 - 13
yamily/__init__.py

@@ -235,9 +235,7 @@ class PersonCollection:
             return self.__children[parent]
         return self.get_children(parent.identifier)
 
-    def get_coparents(
-        self, parent: typing.Union[Person, str]
-    ) -> typing.Iterator[Person]:
+    def get_coparents(self, parent: typing.Union[Person, str]) -> typing.Set[Person]:
         """
         >>> c = PersonCollection()
 
@@ -245,25 +243,25 @@ class PersonCollection:
         >>> alice.father = Person('bob')
         >>> c.add_person(alice)
         Person(alice)
-        >>> list(c.get_coparents('bob'))
-        []
+        >>> c.get_coparents('bob')
+        set()
 
         >>> alice.mother = Person('carol')
         >>> c.add_person(alice)
         Person(alice)
-        >>> list(c.get_coparents('bob'))
-        [Person(carol)]
-        >>> list(c.get_coparents(c['carol']))
-        [Person(bob)]
+        >>> c.get_coparents('bob')
+        {Person(carol)}
+        >>> c.get_coparents(c['carol'])
+        {Person(bob)}
         """
         if isinstance(parent, Person):
-            for coparent in self.get_coparents(parent.identifier):
-                yield coparent
-            return
+            return self.get_coparents(parent.identifier)
+        coparents = set()
         for child in self.get_children(parent):
             for coparent in child.parents:
                 if coparent.identifier != parent:
-                    yield coparent
+                    coparents.add(coparent)
+        return coparents
 
     def __iter__(self) -> "PersonCollection":
         """

+ 85 - 14
yamily/_graphviz.py

@@ -1,34 +1,105 @@
+import typing
+
 import graphviz
 
-from yamily import Person, PersonCollection  # pylint: disable=unused-import; doctest
+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.name = 'Alice'
-    >>> alice.mother = Person('carol')
-    >>> alice.father = Person('bob')
+    >>> alice.mother = carol
+    >>> alice.father = bob
+
+    >>> david = Person('david')
+    >>> david.mother = carol
+    >>> david.father = bob
+
     >>> collection = PersonCollection()
     >>> collection.add_person(alice)
-    Person(alice, Alice)
+    Person(alice)
+    >>> collection.add_person(david)
+    Person(david)
     >>> graph = digraph(collection)
     >>> print(graph.source)
     digraph yamily {
-    	carol [label=unnamed]
-    	bob [label=unnamed]
-    	alice [label=Alice]
-    	carol -> alice
-    	bob -> alice
+    	{
+    		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:
-        graph.node(person.identifier, str(person))
+        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:
-        for parent in [child.mother, child.father]:
-            if parent is not None:
-                graph.edge(parent.identifier, child.identifier)
+        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