diff --git a/dag/ui.go b/dag/ui.go index 18e41dc..e04ef21 100644 --- a/dag/ui.go +++ b/dag/ui.go @@ -131,11 +131,24 @@ func (tm *DAG) ExportDOT() string { sb.WriteString(fmt.Sprintf("digraph \"%s\" {\n", tm.name)) sb.WriteString(" node [shape=box, style=\"rounded,filled\", fillcolor=lightgray, fontname=\"Helvetica\"];\n") sortedNodes := tm.TopologicalSort() + + // Export nodes for _, nodeKey := range sortedNodes { node := tm.nodes[nodeKey] nodeColor := "lightblue" sb.WriteString(fmt.Sprintf(" \"%s\" [label=\"%s\", fillcolor=\"%s\"];\n", node.Key, node.Name, nodeColor)) + + // If the node has a sub-DAG, export it + if subDAG, ok := node.processor.(*DAG); ok && subDAG != nil { + subDAGName := fmt.Sprintf("%s_sub", node.Key) + sb.WriteString(fmt.Sprintf(" subgraph \"%s\" {\n", subDAGName)) + sb.WriteString(fmt.Sprintf(" label=\"%s\"\n", node.Name)) + sb.WriteString(subDAG.ExportDOT()) // Export the sub-DAG + sb.WriteString(" }\n") + } } + + // Export edges for _, nodeKey := range sortedNodes { node := tm.nodes[nodeKey] for _, edge := range node.Edges { @@ -146,23 +159,17 @@ func (tm *DAG) ExportDOT() string { default: edgeStyle = "solid" } + edgeColor := "black" for _, to := range edge.To { - edgeColor := "black" - if edge.Label == "Iterate" { - edgeColor = "blue" - } else if edge.Label == "PASS" { - edgeColor = "green" - } else if edge.Label == "FAIL" { - edgeColor = "red" - } sb.WriteString(fmt.Sprintf(" \"%s\" -> \"%s\" [label=\"%s\", color=\"%s\", style=%s];\n", node.Key, to.Key, edge.Label, edgeColor, edgeStyle)) } } } + + // Handle conditional edges for fromNodeKey, conditions := range tm.conditions { for when, then := range conditions { if toNode, ok := tm.nodes[string(then)]; ok { - sb.WriteString(fmt.Sprintf(" \"%s\" -> \"%s\" [label=\"%s\", color=\"purple\", style=dotted];\n", fromNodeKey, toNode.Key, when)) } } diff --git a/examples/dag.go b/examples/dag.go index 6b16f11..05ee749 100644 --- a/examples/dag.go +++ b/examples/dag.go @@ -51,6 +51,7 @@ func sync() { f := dag.NewDAG("Sample DAG", "sample-dag", mq.WithSyncMode(true), mq.WithNotifyResponse(tasks.NotifyResponse)) setup(f) sendData(f) + fmt.Println(f.SaveSVG("dag.svg")) } func async() {