Write Legend as D2's native legend

src/d2/diagram.py

@@ -21,7 +21,7 @@
     from collections.abc import Mapping
 
     from d2.connection import Connection
-    from d2.shape import LayoutDirection, NearPosition, Shape
+    from d2.shape import LayoutDirection, Shape
     from d2.style import Style
 
 if sys.version_info >= (3, 12):
@@ -47,20 +47,11 @@ def __post_init__(self) -> None:
         bounded("dark_theme_id", self.dark_theme_id, 0)
 
     def statements(self) -> tuple[Field, ...]:
-        entries: tuple[tuple[str, Value | None], ...] = (
-            ("layout-engine", Keyword.of(self.layout_engine)),
-            ("theme-id", self.theme_id),
-            ("dark-theme-id", self.dark_theme_id),
-            ("pad", self.pad),
-            ("sketch", self.sketch),
-            ("center", self.center),
-        )
-        config = properties(entries)
-        if not config:
-            return ()
-        return (
-            Field(Keyword("vars"), body=(Field(Keyword("d2-config"), body=config),)),
-        )
+        return vars_block(self.fields())
+
+    def fields(self) -> tuple[Field, ...]:
+        config = properties(entries_of(self))
+        return (Field(Keyword("d2-config"), body=config),) if config else ()
 
     def lines(self) -> list[str]:
         return emit(self.statements())
@@ -69,20 +60,53 @@ def lines(self) -> list[str]:
 VarValue = str | int | float | bool
 
 
+def vars_block(fields: tuple[Field, ...]) -> tuple[Field, ...]:
+    return (Field(Keyword("vars"), body=fields),) if fields else ()
+
+
+def entries_of(config: Config) -> tuple[tuple[str, Value | None], ...]:
+    return (
+        ("layout-engine", Keyword.of(config.layout_engine)),
+        ("theme-id", config.theme_id),
+        ("dark-theme-id", config.dark_theme_id),
+        ("pad", config.pad),
+        ("sketch", config.sketch),
+        ("center", config.center),
+    )
+
+
+HIDDEN = Field(Keyword("style"), body=(Field(Keyword("opacity"), 0),))
+
+
 @dataclass(frozen=True, kw_only=True)
 class Legend:
+    """D2's legend, drawn as a key next to the diagram.
+
+    Connection endpoints that are not legend shapes are hidden, so an entry
+    for a connection shows only its line and label.
+    """
+
     shapes: list[Shape] = field(default_factory=list)
     connections: list[Connection] = field(default_factory=list)
     label: str | None = None
-    near: NearPosition | None = None
 
-    def statements(self) -> tuple[Field, ...]:
+    def fields(self) -> tuple[Field, ...]:
+        declared = {shape.name for shape in self.shapes}
+        endpoints = dict.fromkeys(
+            end
+            for connection in self.connections
+            for end in (connection.shape_1, connection.shape_2)
+            if end not in declared
+        )
         body: tuple[Statement, ...] = (
-            *properties((("near", Keyword.of(self.near)),)),
             *(s for shape in self.shapes for s in shape.statements()),
+            *(Field(Key(end), Label(""), (HIDDEN,)) for end in endpoints),
             *(s for connection in self.connections for s in connection.statements()),
         )
-        return (Field(Keyword("legend"), Label.of(self.label), body),)
+        return (Field(Keyword("d2-legend"), Label.of(self.label), body),)
+
+    def statements(self) -> tuple[Field, ...]:
+        return vars_block(self.fields())
 
     def lines(self) -> list[str]:
         return emit(self.statements())
@@ -176,23 +200,26 @@ def class_statements(self) -> tuple[Field, ...]:
         )
         return (Field(Keyword("classes"), body=body),) if body else ()
 
-    def variable_statements(self) -> tuple[Field, ...]:
-        body = tuple(
+    def variable_fields(self) -> tuple[Field, ...]:
+        return tuple(
             Field(Key(name), Label(value) if isinstance(value, str) else value)
             for name, value in self.variables.items()
         )
-        return (Field(Keyword("vars"), body=body),) if body else ()
 
     def statements(self) -> tuple[Statement, ...]:
         return (
             *(Comment(comment) for comment in self.comments),
-            *(self.vars.statements() if self.vars is not None else ()),
-            *self.variable_statements(),
+            *vars_block(
+                (
+                    *(self.vars.fields() if self.vars is not None else ()),
+                    *self.variable_fields(),
+                    *(self.legend.fields() if self.legend is not None else ()),
+                )
+            ),
             *properties((("direction", Keyword.of(self.direction)),)),
             *self.class_statements(),
             *(s for shape in self.shapes for s in shape.statements()),
             *(s for connection in self.connections for s in connection.statements()),
-            *(self.legend.statements() if self.legend is not None else ()),
             *boards("layers", self.layers),
             *boards("scenarios", self.scenarios),
             *boards("steps", self.steps),

tests/test_d2/test_d2_vars_and_legend.py

@@ -26,8 +26,6 @@ def test_diagram__variables_follow_config() -> None:
         "  d2-config: {",
         "    sketch: true",
         "  }",
-        "}",
-        "vars: {",
         "  x: 1",
         "}",
     ]
@@ -37,22 +35,52 @@ def test_diagram__legend() -> None:
     diagram = d2.Diagram(
         legend=d2.Legend(
             label="Key",
-            near="bottom-center",
             shapes=[d2.Shape(name="k", label="Service")],
-            connections=[d2.Connection("k", "k", "calls")],
+            connections=[d2.Connection("a", "b", "calls")],
         )
     )
 
     assert str(diagram).split("\n") == [
-        "legend: Key {",
-        "  near: bottom-center",
-        "  k: Service",
-        "  k -> k: calls",
+        "vars: {",
+        "  d2-legend: Key {",
+        "    k: Service",
+        '    a: "" {',
+        "      style: {",
+        "        opacity: 0",
+        "      }",
+        "    }",
+        '    b: "" {',
+        "      style: {",
+        "        opacity: 0",
+        "      }",
+        "    }",
+        "    a -> b: calls",
+        "  }",
         "}",
     ]
 
 
 def test_diagram__legend_without_label() -> None:
     diagram = d2.Diagram(legend=d2.Legend(shapes=[d2.Shape(name="k")]))
 
-    assert str(diagram) == "legend: {\n  k\n}"
+    assert str(diagram) == "vars: {\n  d2-legend: {\n    k\n  }\n}"
+
+
+def test_diagram__one_vars_block() -> None:
+    diagram = d2.Diagram(
+        vars=d2.Config(layout_engine="tala"),
+        variables={"x": 1},
+        legend=d2.Legend(shapes=[d2.Shape(name="k")]),
+    )
+
+    assert str(diagram).split("\n") == [
+        "vars: {",
+        "  d2-config: {",
+        "    layout-engine: tala",
+        "  }",
+        "  x: 1",
+        "  d2-legend: {",
+        "    k",
+        "  }",
+        "}",
+    ]

tests/test_d2/test_render.py

@@ -170,3 +170,21 @@ def fake_run(command: list[str], source: d2.Diagram | str) -> bytes:
     d2.render("a")
 
     assert seen == [["--target", "", "--stdout-format", "svg", "-", "-"]]
+
+
+@requires_d2
+def test_render__legend_is_drawn_as_a_legend() -> None:
+    diagram = d2.Diagram(
+        shapes=[d2.Shape(name="api")],
+        legend=d2.Legend(
+            label="Key",
+            connections=[
+                d2.Connection("a", "b", "calls", style=d2.Style(stroke="red"))
+            ],
+        ),
+    )
+
+    svg = d2.render(diagram).decode()
+
+    assert ">Key<" in svg
+    assert "calls" in svg