dag: call into DotNode to get attributes

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
 		}

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
 

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)