Display if an edge was traveled or not and add all edges

codingjoe · codingjoe · commit e85d0a943060 · 2020-08-14T11:47:32.000+02:00
diff --git a/joeflow/models.py b/joeflow/models.py
@@ -210,7 +210,7 @@ def get_graph(cls, color="black"):
             graph.node(name, style=node_style, color=color, fontcolor=color)
 
         for start, end in cls.edges:
-            graph.edge(start.name, end.name)
+            graph.edge(start.name, end.name, color=color)
         return graph
 
     @classmethod
@@ -265,6 +265,8 @@ def get_instance_graph(self):
                 fontcolor="black",
                 peripheries=peripheries,
             )
+            for child in task.child_task_set.exclude(name="override"):
+                graph.edge(task.name, child.name, color="black")
 
         for task in self.task_set.filter(name="override").prefetch_related(
             "parent_task_set", "child_task_set"
diff --git a/joeflow/utils.py b/joeflow/utils.py
@@ -1,4 +1,5 @@
 import types
+from collections import defaultdict
 
 import graphviz as gv
 
@@ -25,18 +26,64 @@ def get_workflow(name):
 
 
 class NoDashDiGraph(gv.Digraph):
-    """Like `.graphviz.Digraph` but removes underscores from labels."""
+    """
+    Like `.graphviz.Digraph` but with unique nodes and edges.
+
+    Nodes and edges are unique and their attributes will be overridden
+    should the same node or edge be added twice. Nodes are unique by name
+    and edges unique by head and tail.
+
+    Underscores are replaced with whitespaces from identifiers.
+    """
 
     def __init__(self, *args, **kwargs):
-        self._edges = []
+        self._nodes = defaultdict(dict)
+        self._edges = defaultdict(dict)
         super().__init__(*args, **kwargs)
 
-    def edge(self, tail_name, head_name, label=None, _attributes=None, **attrs):
-        if not (tail_name, head_name) in self._edges:
-            self._edges.append((tail_name, head_name))
-            super().edge(
-                tail_name, head_name, label=label, _attributes=_attributes, **attrs
-            )
+    def __iter__(self, subgraph=False):
+        """Yield the DOT source code line by line (as graph or subgraph)."""
+        if self.comment:
+            yield self._comment % self.comment
+
+        if subgraph:
+            if self.strict:
+                raise ValueError("subgraphs cannot be strict")
+            head = self._subgraph if self.name else self._subgraph_plain
+        else:
+            head = self._head_strict if self.strict else self._head
+        yield head % (self._quote(self.name) + " " if self.name else "")
+
+        for kw in ("graph", "node", "edge"):
+            attrs = getattr(self, "%s_attr" % kw)
+            if attrs:
+                yield self._attr % (kw, self._attr_list(None, attrs))
+
+        yield from self.body
+
+        for name, attrs in sorted(self._nodes.items()):
+            name = self._quote(name)
+            label = attrs.pop("label", None)
+            _attributes = attrs.pop("_attributes", None)
+            attr_list = self._attr_list(label, attrs, _attributes)
+            yield self._node % (name, attr_list)
+
+        for edge, attrs in sorted(self._edges.items()):
+            head_name, tail_name = edge
+            tail_name = self._quote_edge(tail_name)
+            head_name = self._quote_edge(head_name)
+            label = attrs.pop("label", None)
+            _attributes = attrs.pop("_attributes", None)
+            attr_list = self._attr_list(label, attrs, _attributes)
+            yield self._edge % (head_name, tail_name, attr_list)
+
+        yield self._tail
+
+    def node(self, name, **attrs):
+        self._nodes[name] = attrs
+
+    def edge(self, tail_name, head_name, **attrs):
+        self._edges[(tail_name, head_name)] = attrs
 
     @staticmethod
     def _quote(identifier, *args, **kwargs):
diff --git a/tests/fixtures/simpleworkflow.dot b/tests/fixtures/simpleworkflow.dot
@@ -1,11 +1,11 @@
 digraph {
 	graph [rankdir=LR]
 	node [fillcolor=white fontname="sans-serif" shape=rect style=filled]
-	"save the princess" [color=black fontcolor=black style="filled, rounded"]
 	end [color=black fontcolor=black style=filled]
-	"start view" [color=black fontcolor=black style="filled, rounded"]
+	"save the princess" [color=black fontcolor=black style="filled, rounded"]
 	"start method" [color=black fontcolor=black style=filled]
-	"start view" -> "save the princess"
-	"start method" -> "save the princess"
-	"save the princess" -> end
+	"start view" [color=black fontcolor=black style="filled, rounded"]
+	"save the princess" -> end [color=black]
+	"start method" -> "save the princess" [color=black]
+	"start view" -> "save the princess" [color=black]
 }
diff --git a/tests/fixtures/simpleworkflow_instance.dot b/tests/fixtures/simpleworkflow_instance.dot
@@ -1,13 +1,11 @@
 digraph {
 	graph [rankdir=LR]
 	node [fillcolor=white fontname="sans-serif" shape=rect style=filled]
-	"start method" [color="#888888" fontcolor="#888888" style=filled]
-	"save the princess" [color="#888888" fontcolor="#888888" style="filled, rounded"]
-	"start view" [color="#888888" fontcolor="#888888" style="filled, rounded"]
 	end [color="#888888" fontcolor="#888888" style=filled]
-	"start view" -> "save the princess"
-	"start method" -> "save the princess"
-	"save the princess" -> end
-	"save the princess" [color=black fontcolor=black href="{url}" peripheries=1 style="filled, rounded, bold"]
+	"save the princess" [color=black fontcolor=black href="/simple/save_the_princess/2/" peripheries=1 style="filled, rounded, bold"]
 	"start method" [color=black fontcolor=black peripheries=1 style="filled, bold"]
+	"start view" [color="#888888" fontcolor="#888888" style="filled, rounded"]
+	"save the princess" -> end [color="#888888"]
+	"start method" -> "save the princess" [color=black]
+	"start view" -> "save the princess" [color="#888888"]
 }
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -108,7 +108,8 @@ def test_get_graph(self, fixturedir):
         assert isinstance(graph, Digraph)
         with open(str(fixturedir / "simpleworkflow.dot")) as fp:
             expected_graph = fp.read().splitlines()
-        assert set(graph.body) == set(expected_graph[1:-1])
+        print(str(graph))
+        assert set(graph) == set(expected_graph)
 
     def test_get_graph_svg(self, fixturedir):
         svg = workflows.SimpleWorkflow.get_graph_svg()
@@ -118,10 +119,9 @@ def test_get_instance_graph(self, db, fixturedir):
         wf = workflows.SimpleWorkflow.start_method()
         task_url = wf.task_set.get(name="save_the_princess").get_absolute_url()
         graph = wf.get_instance_graph()
+        print(str(graph))
         with open(str(fixturedir / "simpleworkflow_instance.dot")) as fp:
-            assert set(graph.body) == set(
-                fp.read().replace("{url}", task_url).splitlines()[1:-1]
-            )
+            assert set(graph) == set(fp.read().replace("{url}", task_url).splitlines())
 
     def test_get_instance_graph__override(
         self, db, stub_worker, fixturedir, admin_client
@@ -133,16 +133,17 @@ def test_get_instance_graph__override(
         stub_worker.wait()
         task = wf.task_set.get(name="override")
         graph = wf.get_instance_graph()
+        print(str(graph))
 
         assert (
             f'\t"{task.name} {task.pk}" [peripheries=1 style="filled, rounded, dashed"]'
-            in graph.body
+            in list(graph)
         )
         assert (
             f'\t"save the princess" -> "{task.name} {task.pk}" [style=dashed]'
-            in graph.body
+            in list(graph)
         )
-        assert f'\t"{task.name} {task.pk}" -> end [style=dashed]' in graph.body
+        assert f'\t"{task.name} {task.pk}" -> end [style=dashed]' in list(graph)
 
     def test_get_instance_graph__obsolete(self, db, fixturedir, admin_client):
         workflow = workflows.SimpleWorkflow.objects.create()
@@ -152,13 +153,14 @@ def test_get_instance_graph__obsolete(self, db, fixturedir, admin_client):
         obsolete.parent_task_set.add(start)
         end.parent_task_set.add(obsolete)
         graph = workflow.get_instance_graph()
+        print(str(graph))
 
         assert (
             '\tobsolete [color=black fontcolor=black peripheries=1 style="filled, dashed, bold"]'
-            in graph.body
+            in str(graph)
         )
-        assert '\t"start method" -> obsolete [style=dashed]' in graph.body
-        assert "\tobsolete -> end [style=dashed]" in graph.body
+        assert '\t"start method" -> obsolete [style=dashed]' in list(graph)
+        assert "\tobsolete -> end [style=dashed]" in list(graph)
 
     def test_get_instance_graph_svg(self, db, fixturedir):
         wf = workflows.SimpleWorkflow.start_method()
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -0,0 +1,73 @@
+from joeflow.utils import NoDashDiGraph
+
+
+class TestNoDashDiGraph:
+    def test_node(self):
+        graph = NoDashDiGraph()
+        graph.node("foo", color="blue")
+        assert list(graph) == [
+            "digraph {",
+            "\tfoo [color=blue]",
+            "}",
+        ]
+        graph.node("foo", color="red")
+        assert list(graph) == [
+            "digraph {",
+            "\tfoo [color=red]",
+            "}",
+        ]
+
+    def test_edge(self):
+        graph = NoDashDiGraph()
+        graph.edge("foo", "bar", color="blue")
+        assert list(graph) == [
+            "digraph {",
+            "\tfoo -> bar [color=blue]",
+            "}",
+        ]
+        graph.edge("foo", "bar", color="red")
+        assert list(graph) == [
+            "digraph {",
+            "\tfoo -> bar [color=red]",
+            "}",
+        ]
+
+    def test_iter(self):
+        graph = NoDashDiGraph(node_attr={"style": "filled"})
+        graph.node("foo", color="red")
+        graph.node("bar", color="green")
+        graph.edge("foo", "bar", color="blue")
+        graph.comment = "This is a comment."
+        print(str(graph))
+        assert list(graph.__iter__()) == [
+            "// This is a comment.",
+            "digraph {",
+            "\tnode [style=filled]",
+            "\tbar [color=green]",
+            "\tfoo [color=red]",
+            "\tfoo -> bar [color=blue]",
+            "}",
+        ]
+
+    def test_iter__subgraph(self):
+        graph = NoDashDiGraph(node_attr={"style": "filled"})
+        graph.node("foo", color="red")
+        graph.node("bar", color="green")
+        graph.edge("foo", "bar", color="blue")
+        graph.comment = "This is a comment."
+        print(str(graph))
+        assert list(graph.__iter__(subgraph=True)) == [
+            "// This is a comment.",
+            "{",
+            "\tnode [style=filled]",
+            "\tbar [color=green]",
+            "\tfoo [color=red]",
+            "\tfoo -> bar [color=blue]",
+            "}",
+        ]
+
+    def test_quote(self):
+        assert NoDashDiGraph._quote("foo_bar") == '"foo bar"'
+
+    def test_quote_edge(self):
+        assert NoDashDiGraph._quote_edge("foo_bar") == '"foo bar"'
