diff --git a/dag/dot.go b/dag/dot.go index fc7754c12..2deb7ccd3 100644 --- a/dag/dot.go +++ b/dag/dot.go @@ -85,8 +85,29 @@ func (v *marshalVertex) dot(g *marshalGraph) []byte { if graphName == "" { graphName = "root" } - buf.WriteString(fmt.Sprintf(`"[%s] %s"`, graphName, v.Name)) - writeAttrs(&buf, v.Attrs) + + name := v.Name + attrs := v.Attrs + if v.graphNodeDotter != nil { + node := v.graphNodeDotter.DotNode(name, nil) + if node == nil { + return []byte{} + } + + newAttrs := make(map[string]string) + for k, v := range attrs { + newAttrs[k] = v + } + for k, v := range node.Attrs { + newAttrs[k] = v + } + + name = node.Name + attrs = newAttrs + } + + buf.WriteString(fmt.Sprintf(`"[%s] %s"`, graphName, name)) + writeAttrs(&buf, attrs) buf.WriteByte('\n') return buf.Bytes() @@ -145,7 +166,7 @@ func (g *marshalGraph) writeBody(opts *DotOpts, w *indentWriter) { skip := map[string]bool{} for _, v := range g.Vertices { - if !v.graphNodeDotter { + if v.graphNodeDotter == nil { skip[v.ID] = true continue } diff --git a/dag/marshal.go b/dag/marshal.go index 77bdc6509..16d5dd6dd 100644 --- a/dag/marshal.go +++ b/dag/marshal.go @@ -102,32 +102,24 @@ type marshalVertex struct { Attrs map[string]string `json:",omitempty"` // This is to help transition from the old Dot interfaces. We record if the - // node was a GraphNodeDotter here, so we know if it should be included in the - // dot output - graphNodeDotter bool + // node was a GraphNodeDotter here, so we can call it to get attributes. + graphNodeDotter GraphNodeDotter } func newMarshalVertex(v Vertex) *marshalVertex { + dn, ok := v.(GraphNodeDotter) + if !ok { + dn = nil + } + return &marshalVertex{ ID: marshalVertexID(v), Name: VertexName(v), Attrs: make(map[string]string), - graphNodeDotter: isDotter(v), + graphNodeDotter: dn, } } -func isDotter(v Vertex) bool { - dn, isDotter := v.(GraphNodeDotter) - dotOpts := &DotOpts{ - Verbose: true, - DrawCycles: true, - } - if isDotter && dn.DotNode("fake", dotOpts) == nil { - isDotter = false - } - return isDotter -} - // vertices is a sort.Interface implementation for sorting vertices by ID type vertices []*marshalVertex diff --git a/dag/marshal_test.go b/dag/marshal_test.go index 14daf6424..9201e0208 100644 --- a/dag/marshal_test.go +++ b/dag/marshal_test.go @@ -34,6 +34,27 @@ func TestGraphDot_basic(t *testing.T) { } } +func TestGraphDot_attrs(t *testing.T) { + var g Graph + g.Add(&testGraphNodeDotter{ + Result: &DotNode{ + Name: "foo", + Attrs: map[string]string{"foo": "bar"}, + }, + }) + + actual := strings.TrimSpace(string(g.Dot(nil))) + expected := strings.TrimSpace(testGraphDotAttrsStr) + if actual != expected { + t.Fatalf("bad: %s", actual) + } +} + +type testGraphNodeDotter struct{ Result *DotNode } + +func (n *testGraphNodeDotter) Name() string { return n.Result.Name } +func (n *testGraphNodeDotter) DotNode(string, *DotOpts) *DotNode { return n.Result } + const testGraphDotBasicStr = `digraph { compound = "true" newrank = "true" @@ -50,6 +71,14 @@ const testGraphDotEmptyStr = `digraph { } }` +const testGraphDotAttrsStr = `digraph { + compound = "true" + newrank = "true" + subgraph "root" { + "[root] foo" [foo = "bar"] + } +}` + func TestGraphJSON_empty(t *testing.T) { var g Graph g.Add(1)