Upload layout_engine.py with huggingface_hub

layout_engine.py

@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""Semantic spec -> clean, valid SAP .drawio. The model's job shrinks to emitting
+typed components + zones + edges; THIS engine does deterministic graphviz layout
+and snaps labels to real SAP icon keys. No pixel coords or style keys from the model.
+
+spec = {
+  "title": str,
+  "zones": [{"id": str, "label": str, "children": [node_id,...]}],
+  "nodes": [{"id": str, "label": str, "icon": optional key}],
+  "edges": [{"source": id, "target": id, "label": optional}]
+}
+"""
+import json, re, subprocess, difflib
+from xml.sax.saxutils import quoteattr
+
+_LIB = json.load(open("dataset/style_library.json"))
+_ICON_KEYS = [k for k in _LIB if not k.startswith(("box", "edge"))]
+
+def slug(s):
+    return re.sub(r"[^a-z0-9]+", "_", (s or "").lower()).strip("_")
+
+def snap_icon(label, given=None):
+    """Return a library key for this label (given key wins, else fuzzy match)."""
+    if given and given in _LIB:
+        return given
+    s = slug(label)
+    if s in _LIB:
+        return s
+    m = difflib.get_close_matches(s, _ICON_KEYS, n=1, cutoff=0.72)
+    return m[0] if m else None
+
+def _dot_layout(spec):
+    """Run graphviz dot, return {node_id: (x,y,w,h)} in pixels, top-left origin."""
+    lines = ['digraph G {', 'rankdir=LR; nodesep=0.5; ranksep=0.9;',
+             'node [shape=box, fixedsize=true, width=1.1, height=0.7];']
+    inzone = set()
+    for z in spec.get("zones", []):
+        lines.append(f'subgraph cluster_{slug(z["id"])} {{ label={json.dumps(z.get("label",""))};')
+        for c in z.get("children", []):
+            lines.append(f'"{c}";'); inzone.add(c)
+        lines.append('}')
+    for n in spec["nodes"]:
+        lines.append(f'"{n["id"]}" [label={json.dumps(n["label"])}];')
+    for e in spec.get("edges", []):
+        lines.append(f'"{e["source"]}" -> "{e["target"]}";')
+    lines.append('}')
+    dot = "\n".join(lines)
+    plain = subprocess.run(["dot", "-Tplain"], input=dot, capture_output=True, text=True).stdout
+    scale = 1.0; H = 0.0; pos = {}
+    for ln in plain.splitlines():
+        p = ln.split()
+        if p[0] == "graph":
+            scale = float(p[1]); H = float(p[3])
+        elif p[0] == "node":
+            name = p[1].strip('"')
+            x, y, w, h = map(float, p[2:6])
+            # graphviz: inches, origin bottom-left -> pixels, top-left
+            PX = 72
+            pos[name] = (int((x - w/2)*PX), int((H - y - h/2)*PX), int(w*PX), int(h*PX))
+    return pos
+
+def _grid_layout(spec):
+    """Zone-column layout with SMALLER icons + an OUTER wrapping block. zones =
+    columns L->R inside the outer block; nodes stacked. Returns node pos, zone
+    rects, and the outer-block rect."""
+    NW, NH, GAP, PAD, TITLE, COLGAP, TOP, MARGIN = 96, 54, 30, 18, 30, 46, 64, 30
+    zones = spec.get("zones", [])
+    zoned = set()
+    for z in zones:
+        zoned.update(z["children"])
+    unzoned = [n["id"] for n in spec["nodes"] if n["id"] not in zoned]
+    columns = ([{"id": None, "label": None, "children": unzoned}] if unzoned else []) + \
+              [{"id": z["id"], "label": z.get("label", "Zone"), "children": z["children"]} for z in zones]
+    pos, zrects = {}, []
+    x = TOP + MARGIN
+    maxy = TOP
+    for col in columns:
+        kids = [k for k in col["children"] if any(n["id"] == k for n in spec["nodes"])]
+        if not kids:
+            continue
+        col_w = NW + 2*PAD
+        y = TOP + MARGIN + (TITLE if col["id"] else 0) + PAD
+        for k in kids:
+            pos[k] = (x + PAD, y, NW, NH)
+            y += NH + GAP
+        if col["id"]:
+            zrects.append({"label": col["label"], "rect": [x, TOP + MARGIN, col_w, y - (TOP+MARGIN) - GAP + PAD]})
+        maxy = max(maxy, y)
+        x += col_w + COLGAP
+    outer = [TOP, TOP, x - COLGAP - TOP + MARGIN, maxy - TOP - GAP + PAD + MARGIN]
+    return pos, zrects, outer
+
+def build(spec, out_path="spec_output.drawio"):
+    pos, zrects, outer = _grid_layout(spec)
+    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
+    # outer surrounding block (behind everything)
+    ox, oy, ow, oh = outer
+    cells.append(
+        f'<mxCell id="outer" value={quoteattr(spec.get("title",""))} '
+        f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#fbfdff;strokeColor=#7a9bbe;'
+        f'verticalAlign=top;fontSize=16;fontStyle=1;fontColor=#13334f;dashed=0;" '
+        f'vertex="1" parent="1"><mxGeometry x="{ox}" y="{oy}" width="{ow}" height="{oh}" as="geometry" /></mxCell>')
+    zid = 1000
+    for z in zrects:
+        x0, y0, w0, h0 = z["rect"]; zid += 1
+        cells.append(
+            f'<mxCell id="z{zid}" value={quoteattr(z.get("label") or "")} '
+            f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#eef5fc;strokeColor=#9fc5e8;'
+            f'verticalAlign=top;fontSize=13;fontStyle=1;fontColor=#1a3b5c;" vertex="1" parent="1">'
+            f'<mxGeometry x="{x0}" y="{y0}" width="{w0}" height="{h0}" as="geometry" /></mxCell>')
+    idmap = {}; nid = 1
+    for n in spec["nodes"]:
+        if n["id"] not in pos:
+            continue
+        x, y, w, h = pos[n["id"]]
+        nid += 1; idmap[n["id"]] = f"n{nid}"
+        key = snap_icon(n["label"], n.get("icon"))
+        if key:
+            style = _LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
+        else:
+            style = "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;"
+        cells.append(
+            f'<mxCell id="n{nid}" value={quoteattr(n["label"])} style={quoteattr(style)} '
+            f'vertex="1" parent="1"><mxGeometry x="{x}" y="{y}" width="{w}" height="{h}" '
+            f'as="geometry" /></mxCell>')
+    # edges
+    eid = 5000
+    for e in spec.get("edges", []):
+        if e["source"] not in idmap or e["target"] not in idmap:
+            continue
+        eid += 1
+        lab = e.get("label", "")
+        va = f' value={quoteattr(lab)}' if lab else ''
+        cells.append(
+            f'<mxCell id="e{eid}"{va} style="edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;'
+            f'strokeColor=#5b7a9c;" edge="1" parent="1" source={quoteattr(idmap[e["source"]])} '
+            f'target={quoteattr(idmap[e["target"]])}><mxGeometry relative="1" as="geometry" /></mxCell>')
+    xml = (f'<mxfile host="archgen"><diagram name={quoteattr(spec.get("title","Architecture"))}>'
+           f'<mxGraphModel dx="1200" dy="800" grid="0" guides="1" tooltips="1" connect="1" '
+           f'arrows="1" fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="1100" '
+           f'math="0" shadow="0"><root>' + "".join(cells) +
+           '</root></mxGraphModel></diagram></mxfile>')
+    open(out_path, "w").write(xml)
+    return out_path
+
+# ---------------- recursive nested-block layout (hierarchical spec) ----------------
+# wide cell (gives the label room) + SMALL icon; row pitch leaves space for the
+# bottom-positioned 2-line label so it never overlaps the next icon.
+ICON_W, ICON_H, BPAD, BTITLE = 124, 40, 20, 34
+COLGAP, ROWGAP, LABEL_H = 34, 26, 34
+BGAP = COLGAP  # block-to-block spacing
+ROW_PITCH = ICON_H + LABEL_H + ROWGAP
+
+def _measure(block):
+    cw, ch = [], []
+    for sub in block.get("blocks", []):
+        _measure(sub); cw.append(sub["_w"]); ch.append(sub["_h"])
+    comps = block.get("components", [])
+    if comps:
+        ncol = 1 if len(comps) <= 5 else 2
+        rows = (len(comps) + ncol - 1) // ncol
+        clw = ncol*ICON_W + (ncol-1)*COLGAP
+        clh = rows*ROW_PITCH - ROWGAP            # no trailing gap; last row keeps its label
+        block["_cluster"] = (ncol, clw, clh); cw.append(clw); ch.append(clh)
+    inner_w = (sum(cw) + COLGAP*(len(cw)-1)) if cw else ICON_W
+    inner_h = max(ch) if ch else ICON_H
+    block["_w"] = inner_w + 2*BPAD
+    block["_h"] = inner_h + BTITLE + BPAD
+
+def _place(block, x, y, pos, rects, depth=0):
+    rects.append({"label": block.get("label", ""), "rect": [x, y, block["_w"], block["_h"]], "depth": depth})
+    cx = x + BPAD; cy = y + BTITLE
+    for sub in block.get("blocks", []):
+        _place(sub, cx, cy, pos, rects, depth+1); cx += sub["_w"] + COLGAP
+    comps = block.get("components", [])
+    if comps:
+        ncol = block["_cluster"][0]
+        for i, comp in enumerate(comps):
+            r, c = divmod(i, ncol)
+            pos[comp["id"]] = (cx + c*(ICON_W+COLGAP), cy + r*ROW_PITCH, ICON_W, ICON_H)
+
+def compute_layout(spec):
+    """Return {component_id: (x, y, w, h)} icon-cell positions (no file written)."""
+    return compute_full(spec)[0]
+
+def compute_full(spec):
+    """Return (pos {id:(x,y,w,h)}, block_rects [{label,rect,depth}]) — for streaming/render."""
+    import copy as _copy
+    blocks = _copy.deepcopy(spec.get("blocks", []))
+    for b in blocks: _measure(b)
+    pos, rects = {}, []; x = 50
+    for b in blocks:
+        _place(b, x, 60, pos, rects, 1); x += b["_w"] + 50
+    return pos, rects
+
+def layout_issues(spec, tol=4):
+    """Visual sanity check: flag any two component boxes (icon + label area) that overlap.
+    The engine is overlap-free by construction, so this should always pass."""
+    pos = compute_layout(spec)
+    # each node occupies icon cell + the label strip below it (LABEL_H tall)
+    boxes = {cid: (x, y, w, ICON_H + LABEL_H) for cid, (x, y, w, h) in pos.items()}
+    ids = list(boxes); out = []
+    for i in range(len(ids)):
+        ax, ay, aw, ah = boxes[ids[i]]
+        for j in range(i + 1, len(ids)):
+            bx, by, bw, bh = boxes[ids[j]]
+            if ax < bx+bw-tol and bx < ax+aw-tol and ay < by+bh-tol and by < ay+ah-tol:
+                out.append(f"layout overlap: {ids[i]} & {ids[j]}")
+    return out[:6]
+
+def _fix_img(style):
+    # browsers/mxGraph need ;base64 for base64 data URIs (draw.io tolerates the comma form)
+    return re.sub(r"data:image/(svg\+xml|png|jpeg),", r"data:image/\1;base64,", style)
+
+def stream_cells(spec):
+    """Real draw.io cells (blocks → nodes → edges) in build order, with EXACT styles +
+    geometry — for incremental insertion into an mxGraph (pixel-identical to build_hier)."""
+    pos, rects = compute_full(spec)
+    fills = ["#eef5fc", "#e4eefb", "#dbe8f8"]
+    idlabel = {}
+    def walk(bs):
+        for b in bs:
+            for c in b.get("components", []): idlabel[c["id"]] = c["label"]
+            walk(b.get("blocks", []))
+    walk(spec.get("blocks", []))
+    cells = []
+    for i, z in enumerate(sorted(rects, key=lambda r: r["depth"])):
+        x, y, w, h = z["rect"]; d = z["depth"]
+        fc = fills[min(d-1, len(fills)-1)] if d >= 1 else "#ffffff"
+        cells.append({"id": f"b{i}", "vertex": True, "value": z.get("label") or "",
+                      "x": x, "y": y, "w": w, "h": h,
+                      "style": f"rounded=0;whiteSpace=wrap;html=1;fillColor={fc};strokeColor=#9fc5e8;"
+                               f"verticalAlign=top;fontSize=12;fontStyle=1;fontColor=#1a3b5c;"})
+    nmap = {}
+    for cid, (x, y, w, h) in pos.items():
+        lbl = idlabel.get(cid, ""); key = snap_icon(lbl)
+        # use the EXACT style build_hier uses (draw.io renders the original comma data-URI form)
+        style = (_LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
+                 "imageAlign=center;align=center;whiteSpace=wrap;") if key \
+                else "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;align=center;"
+        nmap[cid] = f"n{cid}"
+        cells.append({"id": nmap[cid], "vertex": True, "value": lbl,
+                      "x": x, "y": y, "w": w, "h": ICON_H, "style": style})
+    for i, e in enumerate(spec.get("connections", [])):
+        if e["source"] in nmap and e["target"] in nmap:
+            cells.append({"id": f"e{i}", "edge": True, "value": e.get("label", ""),
+                          "source": nmap[e["source"]], "target": nmap[e["target"]],
+                          "style": "edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;strokeColor=#42607d;endArrow=block;"})
+    return cells
+
+def build_hier(spec, out_path="hier_output.drawio"):
+    import copy as _copy
+    blocks = _copy.deepcopy(spec.get("blocks", []))   # don't pollute caller's spec with _w/_h
+    for b in blocks: _measure(b)
+    pos, rects = {}, []
+    x = 50
+    maxh = 0
+    placed = []
+    for b in blocks:
+        _place(b, x, 60, pos, rects, 1); x += b["_w"] + 50; maxh = max(maxh, b["_h"])
+    # outer wrapper
+    ow = x - 50 - 50 + 2*30; oh = maxh + 2*30
+    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
+    cells.append(f'<mxCell id="outer" value={quoteattr(spec.get("title",""))} '
+                 f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;'
+                 f'verticalAlign=top;fontSize=17;fontStyle=1;fontColor=#13334f;" vertex="1" parent="1">'
+                 f'<mxGeometry x="20" y="20" width="{max(ow, x-50)}" height="{oh+60}" as="geometry" /></mxCell>')
+    fills = ["#eef5fc", "#e4eefb", "#dbe8f8"]
+    for i, z in enumerate(sorted(rects, key=lambda r: r["depth"])):
+        x0, y0, w0, h0 = z["rect"]
+        fc = fills[min(z["depth"]-1, len(fills)-1)] if z["depth"] >= 1 else "#ffffff"
+        cells.append(f'<mxCell id="blk{i}" value={quoteattr(z.get("label") or "")} '
+                     f'style="rounded=0;whiteSpace=wrap;html=1;fillColor={fc};strokeColor=#9fc5e8;'
+                     f'verticalAlign=top;fontSize=12;fontStyle=1;fontColor=#1a3b5c;" vertex="1" parent="1">'
+                     f'<mxGeometry x="{x0}" y="{y0}" width="{w0}" height="{h0}" as="geometry" /></mxCell>')
+    idmap = {}; nid = 1
+    def walk(b):
+        for sub in b.get("blocks", []): walk(sub)
+        for comp in b.get("components", []):
+            if comp["id"] not in pos: continue
+            nonlocal nid; nid += 1; idmap[comp["id"]] = f"n{nid}"
+            xx, yy, ww, hh = pos[comp["id"]]
+            key = snap_icon(comp["label"], comp.get("icon"))
+            # small icon centered in the wide cell; label wraps below within the full width
+            style = (_LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
+                     "imageAlign=center;labelPosition=center;align=center;whiteSpace=wrap;") if key \
+                    else "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;align=center;"
+            cells.append(f'<mxCell id="n{nid}" value={quoteattr(comp["label"])} style={quoteattr(style)} '
+                         f'vertex="1" parent="1"><mxGeometry x="{xx}" y="{yy}" width="{ww}" height="{hh}" as="geometry" /></mxCell>')
+    for b in blocks: walk(b)
+    eid = 5000
+    for e in spec.get("connections", []):
+        if e["source"] not in idmap or e["target"] not in idmap: continue
+        eid += 1; lab = e.get("label", ""); va = f' value={quoteattr(lab)}' if lab else ''
+        cells.append(f'<mxCell id="e{eid}"{va} style="edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;'
+                     f'strokeColor=#42607d;endArrow=block;" edge="1" parent="1" '
+                     f'source={quoteattr(idmap[e["source"]])} target={quoteattr(idmap[e["target"]])}>'
+                     f'<mxGeometry relative="1" as="geometry" /></mxCell>')
+    xml = (f'<mxfile host="archgen"><diagram name={quoteattr(spec.get("title","Architecture"))}>'
+           f'<mxGraphModel dx="1200" dy="800" grid="0" guides="1" tooltips="1" connect="1" arrows="1" '
+           f'fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="1100" math="0" shadow="0"><root>'
+           + "".join(cells) + '</root></mxGraphModel></diagram></mxfile>')
+    open(out_path, "w").write(xml)
+    return out_path
+
+if __name__ == "__main__":
+    import sys
+    spec = json.load(open(sys.argv[1])) if len(sys.argv) > 1 else json.load(open("demo_spec.json"))
+    fn = build_hier if "blocks" in spec else build
+    p = fn(spec, sys.argv[2] if len(sys.argv) > 2 else "spec_output.drawio")
+    print("wrote", p)