Merge pull request #2405 from alixander/legend

export diagram.Legend

d2ast/keywords.go

@@ -8,7 +8,6 @@ var ReservedKeywords map[string]struct{}
 // Non Style/Holder keywords.
 var SimpleReservedKeywords = map[string]struct{}{
 	"label":          {},
-	"legend-label":   {},
 	"shape":          {},
 	"icon":           {},
 	"constraint":     {},

d2compiler/compile.go

@@ -96,6 +96,8 @@ func (c *compiler) compileBoard(g *d2graph.Graph, ir *d2ir.Map) *d2graph.Graph {
 	c.validateEdges(g)
 	c.validatePositionsCompatibility(g)
 
+	c.compileLegend(g, ir)
+
 	c.compileBoardsField(g, ir, "layers")
 	c.compileBoardsField(g, ir, "scenarios")
 	c.compileBoardsField(g, ir, "steps")
@@ -110,6 +112,42 @@ func (c *compiler) compileBoard(g *d2graph.Graph, ir *d2ir.Map) *d2graph.Graph {
 	return g
 }
 
+func (c *compiler) compileLegend(g *d2graph.Graph, m *d2ir.Map) {
+	varsField := m.GetField(d2ast.FlatUnquotedString("vars"))
+	if varsField == nil || varsField.Map() == nil {
+		return
+	}
+
+	legendField := varsField.Map().GetField(d2ast.FlatUnquotedString("d2-legend"))
+	if legendField == nil || legendField.Map() == nil {
+		return
+	}
+
+	legendGraph := d2graph.NewGraph()
+
+	c.compileMap(legendGraph.Root, legendField.Map())
+	c.setDefaultShapes(legendGraph)
+
+	objects := make([]*d2graph.Object, 0)
+	for _, obj := range legendGraph.Objects {
+		if obj.Style.Opacity != nil {
+			if opacity, err := strconv.ParseFloat(obj.Style.Opacity.Value, 64); err == nil && opacity == 0 {
+				continue
+			}
+		}
+		objects = append(objects, obj)
+	}
+
+	legend := &d2graph.Legend{
+		Objects: objects,
+		Edges:   legendGraph.Edges,
+	}
+
+	if len(legend.Objects) > 0 || len(legend.Edges) > 0 {
+		g.Legend = legend
+	}
+}
+
 func (c *compiler) compileBoardsField(g *d2graph.Graph, ir *d2ir.Map, fieldName string) {
 	boards := ir.GetField(d2ast.FlatUnquotedString(fieldName))
 	if boards.Map() == nil {
@@ -543,10 +581,6 @@ func (c *compiler) compileReserved(attrs *d2graph.Attributes, f *d2ir.Field) {
 		attrs.Tooltip = &d2graph.Scalar{}
 		attrs.Tooltip.Value = scalar.ScalarString()
 		attrs.Tooltip.MapKey = f.LastPrimaryKey()
-	case "legend-label":
-		attrs.LegendLabel = &d2graph.Scalar{}
-		attrs.LegendLabel.Value = scalar.ScalarString()
-		attrs.LegendLabel.MapKey = f.LastPrimaryKey()
 	case "width":
 		_, err := strconv.Atoi(scalar.ScalarString())
 		if err != nil {

d2compiler/compile_test.go

@@ -720,6 +720,142 @@ x: {
 				}
 			},
 		},
+		{
+			name: "legend",
+
+			text: `
+			vars: {
+  d2-legend: {
+    User: "A person who interacts with the system" {
+      shape: person
+      style: {
+        fill: "#f5f5f5"
+      }
+    }
+
+    Database: "Stores application data" {
+      shape: cylinder
+      style.fill: "#b5d3ff"
+    }
+
+    HiddenShape: "This should not appear in the legend" {
+      style.opacity: 0
+    }
+
+    User -> Database: "Reads data" {
+      style.stroke: "blue"
+    }
+
+    Database -> User: "Returns results" {
+      style.stroke-dash: 5
+    }
+  }
+}
+
+user: User
+db: Database
+user -> db: Uses
+`,
+			assertions: func(t *testing.T, g *d2graph.Graph) {
+				if g.Legend == nil {
+					t.Fatal("Expected Legend to be non-nil")
+					return
+				}
+
+				// 2. Verify the correct objects are in the legend
+				if len(g.Legend.Objects) != 2 {
+					t.Errorf("Expected 2 objects in legend, got %d", len(g.Legend.Objects))
+				}
+
+				// Check for User object
+				hasUser := false
+				hasDatabase := false
+				for _, obj := range g.Legend.Objects {
+					if obj.ID == "User" {
+						hasUser = true
+						if obj.Shape.Value != "person" {
+							t.Errorf("User shape incorrect, expected 'person', got: %s", obj.Shape.Value)
+						}
+					} else if obj.ID == "Database" {
+						hasDatabase = true
+						if obj.Shape.Value != "cylinder" {
+							t.Errorf("Database shape incorrect, expected 'cylinder', got: %s", obj.Shape.Value)
+						}
+					} else if obj.ID == "HiddenShape" {
+						t.Errorf("HiddenShape should not be in legend due to opacity: 0")
+					}
+				}
+
+				if !hasUser {
+					t.Errorf("User object missing from legend")
+				}
+				if !hasDatabase {
+					t.Errorf("Database object missing from legend")
+				}
+
+				// 3. Verify the correct edges are in the legend
+				if len(g.Legend.Edges) != 2 {
+					t.Errorf("Expected 2 edges in legend, got %d", len(g.Legend.Edges))
+				}
+
+				// Check for expected edges
+				hasReadsEdge := false
+				hasReturnsEdge := false
+				for _, edge := range g.Legend.Edges {
+					if edge.Label.Value == "Reads data" {
+						hasReadsEdge = true
+						// Check edge properties
+						if edge.Style.Stroke == nil {
+							t.Errorf("Reads edge stroke is nil")
+						} else if edge.Style.Stroke.Value != "blue" {
+							t.Errorf("Reads edge stroke incorrect, expected 'blue', got: %s", edge.Style.Stroke.Value)
+						}
+					} else if edge.Label.Value == "Returns results" {
+						hasReturnsEdge = true
+						// Check edge properties
+						if edge.Style.StrokeDash == nil {
+							t.Errorf("Returns edge stroke-dash is nil")
+						} else if edge.Style.StrokeDash.Value != "5" {
+							t.Errorf("Returns edge stroke-dash incorrect, expected '5', got: %s", edge.Style.StrokeDash.Value)
+						}
+					} else if edge.Label.Value == "Hidden connection" {
+						t.Errorf("Hidden connection should not be in legend due to opacity: 0")
+					}
+				}
+
+				if !hasReadsEdge {
+					t.Errorf("'Reads data' edge missing from legend")
+				}
+				if !hasReturnsEdge {
+					t.Errorf("'Returns results' edge missing from legend")
+				}
+
+				// 4. Verify the regular diagram content is still there
+				userObj, hasUserObj := g.Root.HasChild([]string{"user"})
+				if !hasUserObj {
+					t.Errorf("Main diagram missing 'user' object")
+				} else if userObj.Label.Value != "User" {
+					t.Errorf("User label incorrect, expected 'User', got: %s", userObj.Label.Value)
+				}
+
+				dbObj, hasDBObj := g.Root.HasChild([]string{"db"})
+				if !hasDBObj {
+					t.Errorf("Main diagram missing 'db' object")
+				} else if dbObj.Label.Value != "Database" {
+					t.Errorf("DB label incorrect, expected 'Database', got: %s", dbObj.Label.Value)
+				}
+
+				// Check the main edge
+				if len(g.Edges) == 0 {
+					t.Errorf("No edges found in main diagram")
+				} else {
+					mainEdge := g.Edges[0]
+					if mainEdge.Label.Value != "Uses" {
+						t.Errorf("Main edge label incorrect, expected 'Uses', got: %s", mainEdge.Label.Value)
+					}
+				}
+			},
+		},
 		{
 			name: "underscore_edge_nested",
 
@@ -5433,31 +5569,13 @@ b -> c
 				assert.Equal(t, "red", g.Edges[0].Style.Stroke.Value)
 			},
 		},
-		{
-			name: "legend-label",
-			run: func(t *testing.T) {
-				g, _ := assertCompile(t, `
-a.legend-label: This is A
-b: {legend-label: This is B}
-a -> b: {
-  legend-label: "This is a->b"
-}
-`, ``)
-				assert.Equal(t, "a", g.Objects[0].ID)
-				assert.Equal(t, "This is A", g.Objects[0].LegendLabel.Value)
-				assert.Equal(t, "b", g.Objects[1].ID)
-				assert.Equal(t, "This is B", g.Objects[1].LegendLabel.Value)
-				assert.Equal(t, "This is a->b", g.Edges[0].LegendLabel.Value)
-			},
-		},
 		{
 			name: "merge-glob-values",
 			run: func(t *testing.T) {
 				assertCompile(t, `
 "a"
 *.style.stroke-width: 2
 *.style.font-size: 14
-
 a.width: 339
 `, ``)
 			},

d2exporter/export.go

@@ -47,6 +47,26 @@ func Export(ctx context.Context, g *d2graph.Graph, fontFamily *d2fonts.FontFamil
 		diagram.Connections[i] = toConnection(g.Edges[i], g.Theme)
 	}
 
+	if g.Legend != nil {
+		legend := &d2target.Legend{}
+
+		if len(g.Legend.Objects) > 0 {
+			legend.Shapes = make([]d2target.Shape, len(g.Legend.Objects))
+			for i, obj := range g.Legend.Objects {
+				legend.Shapes[i] = toShape(obj, g)
+			}
+		}
+
+		if len(g.Legend.Edges) > 0 {
+			legend.Connections = make([]d2target.Connection, len(g.Legend.Edges))
+			for i, edge := range g.Legend.Edges {
+				legend.Connections[i] = toConnection(edge, g.Theme)
+			}
+		}
+
+		diagram.Legend = legend
+	}
+
 	return diagram, nil
 }
 
@@ -243,9 +263,6 @@ func toShape(obj *d2graph.Object, g *d2graph.Graph) d2target.Shape {
 	if obj.Tooltip != nil {
 		shape.Tooltip = obj.Tooltip.Value
 	}
-	if obj.LegendLabel != nil {
-		shape.LegendLabel = obj.LegendLabel.Value
-	}
 	if obj.Style.Animated != nil {
 		shape.Animated, _ = strconv.ParseBool(obj.Style.Animated.Value)
 	}

d2graph/d2graph.go

@@ -49,6 +49,7 @@ type Graph struct {
 	BaseAST *d2ast.Map `json:"-"`
 
 	Root    *Object   `json:"root"`
+	Legend  *Legend   `json:"legend,omitempty"`
 	Edges   []*Edge   `json:"edges"`
 	Objects []*Object `json:"objects"`
 
@@ -67,6 +68,11 @@ type Graph struct {
 	Data map[string]interface{} `json:"data,omitempty"`
 }
 
+type Legend struct {
+	Objects []*Object `json:"objects,omitempty"`
+	Edges   []*Edge   `json:"edges,omitempty"`
+}
+
 func NewGraph() *Graph {
 	d := &Graph{}
 	d.Root = &Object{
@@ -222,8 +228,7 @@ type Attributes struct {
 
 	// These names are attached to the rendered elements in SVG
 	// so that users can target them however they like outside of D2
-	Classes     []string `json:"classes,omitempty"`
-	LegendLabel *Scalar  `json:"legendLabel,omitempty"`
+	Classes []string `json:"classes,omitempty"`
 }
 
 // ApplyTextTransform will alter the `Label.Value` of the current object based

d2target/d2target.go

@@ -86,14 +86,20 @@ type Diagram struct {
 	Shapes      []Shape      `json:"shapes"`
 	Connections []Connection `json:"connections"`
 
-	Root Shape `json:"root"`
+	Root   Shape   `json:"root"`
+	Legend *Legend `json:"legend,omitempty"`
 	// Maybe Icon can be used as a watermark in the root shape
 
 	Layers    []*Diagram `json:"layers,omitempty"`
 	Scenarios []*Diagram `json:"scenarios,omitempty"`
 	Steps     []*Diagram `json:"steps,omitempty"`
 }
 
+type Legend struct {
+	Shapes      []Shape      `json:"shapes,omitempty"`
+	Connections []Connection `json:"connections,omitempty"`
+}
+
 func (d *Diagram) GetBoard(boardPath []string) *Diagram {
 	if len(boardPath) == 0 {
 		return d
@@ -492,7 +498,6 @@ type Shape struct {
 	PrettyLink   string   `json:"prettyLink,omitempty"`
 	Icon         *url.URL `json:"icon"`
 	IconPosition string   `json:"iconPosition"`
-	LegendLabel  string   `json:"legendLabel,omitempty"`
 
 	// Whether the shape should allow shapes behind it to bleed through
 	// Currently just used for sequence diagram groups
@@ -621,7 +626,6 @@ type Connection struct {
 
 	Animated     bool     `json:"animated"`
 	Tooltip      string   `json:"tooltip"`
-	LegendLabel  string   `json:"legendLabel,omitempty"`
 	Icon         *url.URL `json:"icon"`
 	IconPosition string   `json:"iconPosition,omitempty"`