-
-
Notifications
You must be signed in to change notification settings - Fork 95
/
draw.go
145 lines (127 loc) · 3.96 KB
/
draw.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
// Package draw provides functions for visualizing graph structures. At this
// time, draw supports the DOT language which can be interpreted by Graphviz,
// Grappa, and others.
package draw
import (
"fmt"
"io"
"text/template"
"github.com/dominikbraun/graph"
)
// ToDo: This template should be simplified and split into multiple templates.
const dotTemplate = `strict {{.GraphType}} {
{{range $k, $v := .Attributes}}
{{$k}}="{{$v}}";
{{end}}
{{range $s := .Statements}}
"{{.Source}}" {{if .Target}}{{$.EdgeOperator}} "{{.Target}}" [ {{range $k, $v := .EdgeAttributes}}{{$k}}="{{$v}}", {{end}} weight={{.EdgeWeight}} ]{{else}}[ {{range $k, $v := .SourceAttributes}}{{$k}}="{{$v}}", {{end}} weight={{.SourceWeight}} ]{{end}};
{{end}}
}
`
type description struct {
GraphType string
Attributes map[string]string
EdgeOperator string
Statements []statement
}
type statement struct {
Source interface{}
Target interface{}
SourceWeight int
SourceAttributes map[string]string
EdgeWeight int
EdgeAttributes map[string]string
}
// DOT renders the given graph structure in DOT language into an io.Writer, for
// example a file. The generated output can be passed to Graphviz or other
// visualization tools supporting DOT.
//
// The following example renders a directed graph into a file my-graph.gv:
//
// g := graph.New(graph.IntHash, graph.Directed())
//
// _ = g.AddVertex(1)
// _ = g.AddVertex(2)
// _ = g.AddVertex(3, graph.VertexAttribute("style", "filled"), graph.VertexAttribute("fillcolor", "red"))
//
// _ = g.AddEdge(1, 2, graph.EdgeWeight(10), graph.EdgeAttribute("color", "red"))
// _ = g.AddEdge(1, 3)
//
// file, _ := os.Create("./my-graph.gv")
// _ = draw.DOT(g, file)
//
// To generate an SVG from the created file using Graphviz, use a command such
// as the following:
//
// dot -Tsvg -O my-graph.gv
//
// Another possibility is to use os.Stdout as an io.Writer, print the DOT output
// to stdout, and pipe it as follows:
//
// go run main.go | dot -Tsvg > output.svg
//
// DOT also accepts the [GraphAttribute] functional option, which can be used to
// add global attributes when rendering the graph:
//
// _ = draw.DOT(g, file, draw.GraphAttribute("label", "my-graph"))
func DOT[K comparable, T any](g graph.Graph[K, T], w io.Writer, options ...func(*description)) error {
desc, err := generateDOT(g, options...)
if err != nil {
return fmt.Errorf("failed to generate DOT description: %w", err)
}
return renderDOT(w, desc)
}
// GraphAttribute is a functional option for the [DOT] method.
func GraphAttribute(key, value string) func(*description) {
return func(d *description) {
d.Attributes[key] = value
}
}
func generateDOT[K comparable, T any](g graph.Graph[K, T], options ...func(*description)) (description, error) {
desc := description{
GraphType: "graph",
Attributes: make(map[string]string),
EdgeOperator: "--",
Statements: make([]statement, 0),
}
for _, option := range options {
option(&desc)
}
if g.Traits().IsDirected {
desc.GraphType = "digraph"
desc.EdgeOperator = "->"
}
adjacencyMap, err := g.AdjacencyMap()
if err != nil {
return desc, err
}
for vertex, adjacencies := range adjacencyMap {
_, sourceProperties, err := g.VertexWithProperties(vertex)
if err != nil {
return desc, err
}
stmt := statement{
Source: vertex,
SourceWeight: sourceProperties.Weight,
SourceAttributes: sourceProperties.Attributes,
}
desc.Statements = append(desc.Statements, stmt)
for adjacency, edge := range adjacencies {
stmt := statement{
Source: vertex,
Target: adjacency,
EdgeWeight: edge.Properties.Weight,
EdgeAttributes: edge.Properties.Attributes,
}
desc.Statements = append(desc.Statements, stmt)
}
}
return desc, nil
}
func renderDOT(w io.Writer, d description) error {
tpl, err := template.New("dotTemplate").Parse(dotTemplate)
if err != nil {
return fmt.Errorf("failed to parse template: %w", err)
}
return tpl.Execute(w, d)
}