Add legends to diagrams

Author: Jelle Peters

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/DiagramGenerator.kt

@@ -8,6 +8,7 @@ import com.structurizr.view.View
 import net.sourceforge.plantuml.FileFormat
 import net.sourceforge.plantuml.FileFormatOption
 import net.sourceforge.plantuml.SourceStringReader
+import nl.avisi.structurizr.site.generatr.site.model.DiagramSvgs
 import java.io.ByteArrayOutputStream
 import java.io.File
 import java.net.URI
@@ -23,11 +24,27 @@ fun generateDiagrams(workspace: Workspace, exportDir: File) {
     plantUMLDiagrams.parallelStream()
         .forEach { diagram ->
             val plantUMLFile = File(pumlDir, "${diagram.key}.puml")
-            if (!plantUMLFile.exists() || plantUMLFile.readText() != diagram.definition) {
+            val plantUMLLegendFile = File(pumlDir, "${diagram.key}.legend.puml")
+
+            val diagramNeedsUpdating = !plantUMLFile.exists() || plantUMLFile.readText() != diagram.definition
+            val legendNeedsUpdating = !plantUMLLegendFile.exists() || plantUMLLegendFile.readText() != diagram.legend?.definition
+
+            if (diagramNeedsUpdating || legendNeedsUpdating) {
                 println("${diagram.key}...")
-                saveAsSvg(diagram, svgDir)
-                saveAsPng(diagram, pngDir)
-                saveAsPUML(diagram, plantUMLFile)
+                if (diagramNeedsUpdating) {
+                    saveAsSvg(diagram.withCachedIncludes().definition, svgDir, diagram.key)
+                    saveAsPng(diagram.withCachedIncludes().definition, pngDir, diagram.key)
+                    saveAsPUML(diagram.definition, plantUMLFile)
+                }
+                if (legendNeedsUpdating) {
+                    val diagramLegendDefinition = diagram.legend?.definition
+
+                    if (diagramLegendDefinition != null) {
+                        saveAsSvg(diagramLegendDefinition, svgDir,diagram.key + ".legend")
+                        saveAsPng(diagramLegendDefinition, pngDir, diagram.key + ".legend")
+                        saveAsPUML(diagramLegendDefinition, plantUMLLegendFile)
+                    }
+                }
             } else {
                 println("${diagram.key} UP-TO-DATE")
             }
@@ -38,18 +55,23 @@ fun generateDiagramWithElementLinks(
     workspace: Workspace,
     view: View,
     url: String,
-    diagramCache: ConcurrentHashMap<String, String>
-): String {
+    diagramCache: ConcurrentHashMap<String, DiagramSvgs>
+): DiagramSvgs {
     val diagram = generatePlantUMLDiagramWithElementLinks(workspace, view, url)
 
     val name = "${diagram.key}-${view.key}"
-    return diagramCache.getOrPut(name) {
-        val reader = SourceStringReader(diagram.withCachedIncludes().definition)
-        val stream = ByteArrayOutputStream()
+    return diagramCache.getOrPut(name) { DiagramSvgs(
+        convertDefinitionToSvg(diagram.withCachedIncludes().definition),
+        diagram.legend?.definition?.let { convertDefinitionToSvg(it) }
+    ) }
+}
 
-        reader.outputImage(stream, FileFormatOption(FileFormat.SVG, false))
-        stream.toString(Charsets.UTF_8)
-    }
+private fun convertDefinitionToSvg(definition: String): String {
+    val reader = SourceStringReader(definition)
+    val stream = ByteArrayOutputStream()
+
+    reader.outputImage(stream, FileFormatOption(FileFormat.SVG, false))
+    return stream.toString(Charsets.UTF_8)
 }
 
 private fun generatePlantUMLDiagrams(workspace: Workspace): Collection<Diagram> {
@@ -58,22 +80,22 @@ private fun generatePlantUMLDiagrams(workspace: Workspace): Collection<Diagram>
     return plantUMLExporter.export()
 }
 
-private fun saveAsPUML(diagram: Diagram, plantUMLFile: File) {
-    plantUMLFile.writeText(diagram.definition)
+private fun saveAsPUML(definition: String, plantUMLFile: File) {
+    plantUMLFile.writeText(definition)
 }
 
-private fun saveAsSvg(diagram: Diagram, svgDir: File, name: String = diagram.key) {
-    val reader = SourceStringReader(diagram.withCachedIncludes().definition)
+private fun saveAsSvg(definition: String, svgDir: File, name: String) {
+    val reader = SourceStringReader(definition)
     val svgFile = File(svgDir, "$name.svg")
 
     svgFile.outputStream().use {
         reader.outputImage(it, FileFormatOption(FileFormat.SVG, false))
     }
 }
 
-private fun saveAsPng(diagram: Diagram, pngDir: File) {
-    val reader = SourceStringReader(diagram.withCachedIncludes().definition)
-    val pngFile = File(pngDir, "${diagram.key}.png")
+private fun saveAsPng(definition: String, pngDir: File, name: String) {
+    val reader = SourceStringReader(definition)
+    val pngFile = File(pngDir, "$name.png")
 
     pngFile.outputStream().use {
         reader.outputImage(it)

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/GeneratorContext.kt

@@ -1,12 +1,13 @@
 package nl.avisi.structurizr.site.generatr.site
 
 import com.structurizr.Workspace
+import nl.avisi.structurizr.site.generatr.site.model.DiagramSvgs
 
 data class GeneratorContext(
     val version: String,
     val workspace: Workspace,
     val branches: List<String>,
     val currentBranch: String,
     val serving: Boolean,
-    val svgFactory: (key: String, url: String) -> String?
+    val svgFactory: (key: String, url: String) -> DiagramSvgs?
 )

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/SiteGenerator.kt

@@ -73,7 +73,7 @@ fun generateSite(
     serving: Boolean = false
 ) {
     val generatorContext = GeneratorContext(version, workspace, branches, currentBranch, serving) { key, url ->
-        val diagramCache = ConcurrentHashMap<String, String>()
+        val diagramCache = ConcurrentHashMap<String, DiagramSvgs>()
         workspace.views.views.singleOrNull { view -> view.key == key }
             ?.let { generateDiagramWithElementLinks(workspace, it, url, diagramCache) }
     }

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/model/DiagramSvgs.kt

@@ -0,0 +1,6 @@
+package nl.avisi.structurizr.site.generatr.site.model
+
+data class DiagramSvgs(
+    val svg: String,
+    val legendSvg: String?,
+)

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/model/DiagramViewModel.kt

@@ -10,10 +10,11 @@ data class DiagramViewModel(
     val diagramWidthInPixels: Int?,
     val svgLocation: ImageViewModel,
     val pngLocation: ImageViewModel,
-    val pumlLocation: ImageViewModel
+    val pumlLocation: ImageViewModel,
+    val legend: LegendViewModel?,
 ) {
     companion object {
-        fun forView(pageViewModel: PageViewModel, view: View, svgFactory: (key: String, url: String) -> String?) =
+        fun forView(pageViewModel: PageViewModel, view: View, svgFactory: (key: String, url: String) -> DiagramSvgs?) =
             forView(pageViewModel, view.key, view.name, view.title, view.description.ifBlank { null }, svgFactory)
 
         fun forView(
@@ -22,9 +23,11 @@ data class DiagramViewModel(
             name: String,
             title: String?,
             description: String?,
-            svgFactory: (key: String, url: String) -> String?
+            svgFactory: (key: String, url: String) -> DiagramSvgs?
         ): DiagramViewModel {
-            val svg = svgFactory(key, pageViewModel.url)
+            val (svg, legendSvg) = svgFactory(key, pageViewModel.url)?.let {
+                it.svg to it.legendSvg
+            } ?: (null to null)
             return DiagramViewModel(
                 key,
                 title ?: name,
@@ -33,7 +36,14 @@ data class DiagramViewModel(
                 extractDiagramWidthInPixels(svg),
                 ImageViewModel(pageViewModel, "/svg/$key.svg"),
                 ImageViewModel(pageViewModel, "/png/$key.png"),
-                ImageViewModel(pageViewModel, "/puml/$key.puml")
+                ImageViewModel(pageViewModel, "/puml/$key.puml"),
+                legendSvg?.let { LegendViewModel(
+                    legendSvg,
+                    extractDiagramWidthInPixels(legendSvg),
+                    ImageViewModel(pageViewModel, "/svg/$key.legend.svg"),
+                    ImageViewModel(pageViewModel, "/png/$key.legend.png"),
+                    ImageViewModel(pageViewModel, "/puml/$key.legend.puml"),
+                ) }
             )
         }
 

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/model/LegendViewModel.kt

@@ -0,0 +1,9 @@
+package nl.avisi.structurizr.site.generatr.site.model
+
+data class LegendViewModel(
+    val svg: String?,
+    val widthInPixels: Int?,
+    val svgLocation: ImageViewModel,
+    val pngLocation: ImageViewModel,
+    val pumlLocation: ImageViewModel,
+)

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/model/ToHtml.kt

@@ -1,6 +1,7 @@
 package nl.avisi.structurizr.site.generatr.site.model
 
 import com.structurizr.documentation.Format
+import com.structurizr.export.Diagram
 import com.vladsch.flexmark.ast.FencedCodeBlock
 import com.vladsch.flexmark.html.HtmlRenderer
 import com.vladsch.flexmark.html.HtmlWriter
@@ -29,7 +30,7 @@ fun toHtml(
     pageViewModel: PageViewModel,
     content: String,
     format: Format,
-    svgFactory: (key: String, url: String) -> String?
+    svgFactory: (key: String, url: String) -> DiagramSvgs?
 ): String = when (format) {
     Format.Markdown -> markdownToHtml(pageViewModel, content, svgFactory)
     Format.AsciiDoc -> asciidocToHtml(pageViewModel, content, svgFactory)
@@ -38,7 +39,7 @@ fun toHtml(
 private fun markdownToHtml(
     pageViewModel: PageViewModel,
     markdown: String,
-    svgFactory: (key: String, url: String) -> String?
+    svgFactory: (key: String, url: String) -> DiagramSvgs?
 ): String {
     val flexmarkConfig = pageViewModel.flexmarkConfig
     val options = flexmarkConfig.flexmarkOptions
@@ -77,7 +78,7 @@ private class FencedCodeBlockRenderer : NodeRenderer {
 private fun asciidocToHtml(
     pageViewModel: PageViewModel,
     asciidoc: String,
-    svgFactory: (key: String, url: String) -> String?
+    svgFactory: (key: String, url: String) -> DiagramSvgs?
 ): String {
     val options = Options.builder()
         .safe(SafeMode.SERVER)
@@ -159,7 +160,7 @@ private class CustomLinkResolver(private val pageViewModel: PageViewModel) : Lin
 
 fun Element.transformEmbeddedDiagramElements(
     pageViewModel: PageViewModel,
-    svgFactory: (key: String, url: String) -> String?
+    svgFactory: (key: String, url: String) -> DiagramSvgs?
 ) = this.allElements
     .toList()
     .filter { it.tag().name == "img" && it.attr("src").startsWith(embedPrefix) }

src/main/kotlin/nl/avisi/structurizr/site/generatr/site/views/Diagram.kt

@@ -35,6 +35,15 @@ fun FlowContent.diagram(viewModel: DiagramViewModel) {
                 +"|"
                 a(href = viewModel.pumlLocation.relativeHref, target = "_blank") { +"puml" }
                 +"]"
+                if (viewModel.legend != null) {
+                    +" Legend ["
+                    a(href = viewModel.legend.svgLocation.relativeHref, target = "_blank") { +"svg" }
+                    +"|"
+                    a(href = viewModel.legend.pngLocation.relativeHref, target = "_blank") { +"png" }
+                    +"|"
+                    a(href = viewModel.legend.pumlLocation.relativeHref, target = "_blank") { +"puml" }
+                    +"]"
+                }
             }
         }
     } else

src/test/kotlin/nl/avisi/structurizr/site/generatr/site/StructurizrUtilitiesTest.kt

@@ -8,7 +8,7 @@ import kotlin.test.Test
 
 class StructurizrUtilitiesTest {
 
-    private val svgFactory = { _: String, _: String -> "" }
+    private val svgFactory = { _: String, _: String -> "" to "" }
 
     private fun generatorContext(
         workspaceName: String = "Workspace name",

src/test/kotlin/nl/avisi/structurizr/site/generatr/site/model/AsciidocToHtmlTest.kt

@@ -211,6 +211,11 @@ class AsciidocToHtmlTest : ViewModelTest() {
                 <div>
                  <svg viewBox="0 0 800 900"></svg>
                 </div>
+                <div style="width: min(100%, 800px);">
+                 <div>
+                  <svg viewBox="0 0 800 900"></svg>
+                 </div>
+                </div>
                 <figcaption>
                  <a onclick="openSvgModal('SystemLandscape-modal', 'SystemLandscape-svg')">System Landscape Diagram</a>
                 </figcaption>
@@ -224,6 +229,8 @@ class AsciidocToHtmlTest : ViewModelTest() {
                   </div>
                   <div class="has-text-centered">
                    System Landscape Diagram [<a href="svg/SystemLandscape.svg" target="_blank">svg</a>|<a href="png/SystemLandscape.png" target="_blank">png</a>|<a href="puml/SystemLandscape.puml" target="_blank">puml</a>]
+                   <br>
+                   Legend [<a href="svg/SystemLandscape.legend.svg" target="_blank">svg</a>|<a href="png/SystemLandscape.legend.png" target="_blank">png</a>|<a href="svg/SystemLandscape.legend.svg" target="_blank">puml</a>]
                   </div>
                  </div>
                 </div>