Refactor visualize.py to use jinja2

shell.nix

@@ -6,5 +6,6 @@ mkShell {
 		d2
 		dnscontrol
 		opentofu
+		python3Packages.jinja2
 	];
 }

templates/infrastructure.d2.j2

@@ -0,0 +1,77 @@
+vars: {
+  d2-config: {
+    layout-engine: elk
+    theme-id: 101
+  }
+}
+
+external: {
+    scaleway: {
+        s3
+    }
+
+    restic: {
+        icon: https://cdn.jsdelivr.net/gh/selfhst/icons/webp/restic.webp
+    }
+}
+
+{% for host in hosts %}
+{{ host.key }}: {
+}
+{% endfor %}{# host #}
+
+{% for svc in svcs %}
+{{ svc.key }}: {
+    label: {{ svc.label }}
+    label.near: top-left
+    icon: {{ svc.icon }}
+    icon.near: top-right
+}
+
+{% for backup in svc.backup or [] %}
+{{ svc.key }} -> external.restic.{{ svc.host }}: backup.name {
+    style: {
+        stroke: "#0f0"
+        stroke-dash: 3
+    }
+}
+{% endfor %}{# backup #}
+
+{% if svc.monitoring %}
+{{ monitoring_key }} -> {{ svc.key }}: {
+    style.stroke: "#1E9025"
+}
+{% endif %}
+
+{% if svc.database %}
+{{ svc.key }} -> {{ db_key }}: {
+    style.stroke: "#336791"
+}
+{{ db_key }}.{{ svc.name }}
+{% endif %}
+
+{% if svc.auth %}
+{{ svc.key }} -> {{ auth_key }}: {
+    style.stroke: "#FD4B2D"
+}
+{{ auth_key }}.{{ svc.name }}
+{% endif %}
+
+{% if svc.s3 %}
+{{ svc.key }} -> external.scaleway.s3: {
+    style: {
+        stroke: "#110B1E"
+        stroke-dash: 3
+    }
+}
+external.scaleway.s3.{{ svc.name }}
+{% endif %}
+
+{% endfor %}{# svc #}
+
+{% for svc in grid_svcs %}
+{{ svc }}: {
+    grid-columns: 3
+    grid-gap: 0
+}
+{% endfor %}

visualize.py

@@ -3,6 +3,8 @@
 import json
 import sys
 
+import jinja2
+
 icon_overrides = {
     "acme_dns": "lets-encrypt",
     "extra_services": None,
@@ -24,41 +26,6 @@ icon_format = {
     "tiny-tiny-rss": "webp",
 }
 
-template_host = """
-'serguzim.net'.{host}: {{
-    label: {host}
-}}
-"""
-
-template_service = """
-{key}: {{
-    label: {label}
-    label.near: top-left
-    icon: {icon}
-    icon.near: top-right
-}}
-"""
-
-default_d2 = """
-vars: {
-  d2-config: {
-    layout-engine: elk
-    # Terminal theme code
-    theme-id: 101
-  }
-}
-
-external: {
-    scaleway: {
-        s3
-    }
-
-    restic: {
-        icon: https://cdn.jsdelivr.net/gh/selfhst/icons/webp/restic.webp
-    }
-}
-"""
-
 def get_icon(svc):
     svc = icon_overrides.get(svc, svc) or 'docker'
     fmt = icon_format.get(svc, 'svg')
@@ -66,6 +33,9 @@ def get_icon(svc):
 
 def service_key(svc, data):
     return f"'serguzim.net'.{data['host']}.{svc}"
+def host_key(host):
+    return f"'serguzim.net'.{host}"
+
 
 def service_key_find(svc_name, services):
     for svc, data in services.items():
@@ -75,8 +45,11 @@ def service_key_find(svc_name, services):
 
 def parse_hosts(hosts):
     result = []
-    for host in hosts.keys():
-        result.append(template_host.format(host=host))
+    for host, data in hosts.items():
+        result.append({
+            'key': host_key(host),
+            'name': host,
+        })
     return result
 
 def parse_services(services):
@@ -85,42 +58,23 @@ def parse_services(services):
     postgresql_key = service_key_find("postgresql", services)
     authentik_key = service_key_find("authentik", services)
 
-    result.append(f"{postgresql_key}.grid-columns: 3")
-    result.append(f"{postgresql_key}.grid-gap: 0")
-    result.append(f"{authentik_key}.grid-columns: 3")
-    result.append(f"{authentik_key}.grid-gap: 0")
-
     for svc, data in services.items():
         svc_key = service_key(svc, data)
 
         domains = []
-        if data.get("dns"):
-            domains = []
-            for dns in data["dns"]:
-                domain = ""
-                if dns.get("target") != "@":
-                    domain += f"{dns["target"]}."
-                domain += dns['domain']
-                domains.append(f"- {domain}")
+        for dns in data.get("dns") or []:
+            domain = ""
+            if dns.get("target") != "@":
+                domain += f"{dns["target"]}."
+            domain += dns['domain']
+            domains.append(f"- {domain}")
 
-        result.append(template_service.format(
-            key=svc_key,
-            label="\\n".join([svc] + domains),
-            icon=get_icon(svc)
-        ))
+        data['key'] = svc_key
+        data['label'] = "\\n".join([svc] + domains)
+        data['icon'] = get_icon(svc)
 
-        for backup in data.get("backup") or []:
-            result.append(f'({svc_key} -> external.restic.{data['host']}).style.stroke: "#0f0"')
+        result.append(data)
 
-        if data.get("database"):
-            result.append(f"({svc_key} -> {postgresql_key}).style.stroke: '#00f'")
-            result.append(f"{postgresql_key}.{svc}")
-        if data.get("auth"):
-            result.append(f"({svc_key} -> {authentik_key}).style.stroke: '#FD4B2D'")
-            result.append(f"{authentik_key}.{svc}")
-        if data.get("s3"):
-            result.append(f"({svc_key} -> external.scaleway.s3).style.stroke: '#bbb'")
-            result.append(f"external.scaleway.s3.{svc}")
     return result
 
 if __name__ == '__main__':
@@ -128,11 +82,20 @@ if __name__ == '__main__':
     services = []
 
     data = json.loads(sys.stdin.read())
-    
-    hosts = parse_hosts(data["hosts"])
-    services = parse_services(data["services"])
 
-    print("\n".join(
-        [default_d2]
-            + hosts
-            + services))
+    db_key = service_key_find("postgresql", data["services"])
+    auth_key = service_key_find("authentik", data["services"])
+    monitoring_key = service_key_find("gatus", data["services"])
+
+    jinja_loader = jinja2.FileSystemLoader(searchpath="./templates")
+    jinja_env = jinja2.Environment(loader=jinja_loader)
+    template = jinja_env.get_template("infrastructure.d2.j2")
+    text = template.render(
+        grid_svcs=[db_key, auth_key],
+        svcs=parse_services(data["services"]),
+        hosts=parse_hosts(data["hosts"]),
+        db_key=db_key,
+        auth_key=auth_key,
+        monitoring_key=monitoring_key,
+    )
+    print(text)