From c1eae406e93a1ace28e283e745dc8c381659fb9b Mon Sep 17 00:00:00 2001 From: robby Date: Fri, 23 Oct 2026 15:23:35 -0400 Subject: [PATCH] Write Legend as D2's native legend under vars.d2-legend Legend was written as a top-level 'legend' key, which D2 draws as an ordinary container. It now goes under vars.d2-legend with its label as the title, so D2 draws a real legend. Connection endpoints that are not legend shapes are hidden, so an entry for a connection shows only its line and label. The diagram writes one vars block for d2-config, variables and the legend. Legend.near is removed: D2 places the legend itself and has no position for it. Co-Authored-By: Claude Opus 5.5 --- src/d2/diagram.py | 77 ++++++++++++++++-------- tests/test_d2/test_d2_vars_and_legend.py | 46 +++++++++++--- tests/test_d2/test_render.py | 18 ++++++ 3 files changed, 107 insertions(+), 34 deletions(-) diff --git a/src/d2/diagram.py b/src/d2/diagram.py index 74aece2..f0b8fc8 100644 --- a/src/d2/diagram.py +++ b/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), diff --git a/tests/test_d2/test_d2_vars_and_legend.py b/tests/test_d2/test_d2_vars_and_legend.py index c072381..78fe6c7 100644 --- a/tests/test_d2/test_d2_vars_and_legend.py +++ b/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,17 +35,27 @@ 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", + " }", "}", ] @@ -55,4 +63,24 @@ def test_diagram__legend() -> None: 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", + " }", + "}", + ] diff --git a/tests/test_d2/test_render.py b/tests/test_d2/test_render.py index 7d0b308..011ce0d 100644 --- a/tests/test_d2/test_render.py +++ b/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