extract_blocks.py

#!/usr/bin/env python3
"""Hierarchical extractor: real .drawio -> nested block tree.
blocks(containers) nest by geometric containment; components(icons) are leaves;
connections are edges remapped to nearest component. Robust to parent cycles."""
import re, json, html, math
from drawio_engine import _models, signature
sig2key = json.load(open("dataset/sig2key.json"))

def _label(v):
    if not v: return ""
    return re.sub(r"\s+", " ", re.sub(r"<[^>]+>", " ", html.unescape(v))).strip()
JUNK = re.compile(r"^(l\d|level\s*\d|diagram.*|page-?\d*|\d+|[a-z])$", re.I)
PROTO = re.compile(r"^(soap|rest|rfc|http|https|as2|as4|idoc|edi|odata|tcp|ssl|mqtt|amqp|json|xml)[\s/]*", re.I)
def junk(l): return (not l) or len(l) < 3 or JUNK.match(l) or len(l) > 60
def proto(l): return bool(l) and bool(PROTO.match(l)) and len(l) < 14

def extract_hier(path, title=None, max_pages=1):
    out = []
    for pi, (pname, gm) in enumerate(_models(path)):
        if pi >= max_pages: break
        root = gm.find("root")
        if root is None: continue
        cells = {c.get("id"): c for c in root.iter("mxCell") if c.get("id")}
        geo, parent = {}, {}
        for cid, c in cells.items():
            parent[cid] = c.get("parent"); g = c.find("mxGeometry")
            if g is not None:
                try: geo[cid] = [float(g.get(k,0) or 0) for k in ("x","y","width","height")]
                except: geo[cid] = [0,0,80,40]
        def absr(cid, d=0):
            if cid not in geo: return [0,0,80,40]
            x,y,w,h = geo[cid]; p = parent.get(cid)
            if d < 40 and p in geo and p not in ("0","1"):
                px,py,_,_ = absr(p,d+1); x+=px; y+=py
            return [x,y,w,h]
        comps, blocks, texts = {}, {}, []
        for cid, c in cells.items():
            if c.get("vertex") != "1": continue
            st = c.get("style","") or ""; lab = _label(c.get("value")); r = absr(cid)
            if "image=data:" in st and not junk(lab):
                comps[cid] = {"label": lab, "icon": sig2key.get(signature(st)), "rect": r}
            elif r[2] >= 158 and r[3] >= 88:   # labeled grouping container (e.g. 170x102)
                blocks[cid] = {"label": lab if not junk(lab) else "", "rect": r}
            elif "image=data:" not in st and lab and not junk(lab) and r[3] < 55:
                texts.append({"label": lab, "rect": r})
        if len(comps) < 2: continue
        def contains(a, b):  # a contains b's center
            ax,ay,aw,ah = a; cx,cy = b[0]+b[2]/2, b[1]+b[3]/2
            return ax <= cx <= ax+aw and ay <= cy <= ay+ah
        def area(r): return r[2]*r[3]
        # reject protocol-like block labels, then borrow a real title from text near the top
        for b in blocks.values():
            if proto(b["label"]): b["label"] = ""
            if b["label"]: continue
            bx,by,bw,bh = b["rect"]; best, by2 = None, 1e18
            for t in texts:
                if proto(t["label"]): continue
                cx,cy = t["rect"][0]+t["rect"][2]/2, t["rect"][1]+t["rect"][3]/2
                if bx<=cx<=bx+bw and by-8<=cy<=by+bh*0.35 and cy<by2: by2,best=cy,t["label"]
            if best: b["label"] = best
        # parent block = smallest block strictly containing it
        def parent_block(cid, rect, is_block):
            best, ba = None, 1e18
            for bid, b in blocks.items():
                if bid == cid: continue
                if contains(b["rect"], rect) and area(b["rect"]) > area(rect)*1.02:
                    if area(b["rect"]) < ba: ba, best = area(b["rect"]), bid
            return best
        bparent = {bid: parent_block(bid, b["rect"], True) for bid, b in blocks.items()}
        cparent = {cid: parent_block(cid, c["rect"], False) for cid, c in comps.items()}
        cid_map = {cid: f"c{i}" for i, cid in enumerate(comps)}
        def build_block(bid):
            node = {"id": "b"+str(list(blocks).index(bid)),
                    "label": blocks[bid]["label"] or "", "blocks": [], "components": []}
            for sub, p in bparent.items():
                if p == bid: node["blocks"].append(build_block(sub))
            for c, p in cparent.items():
                if p == bid: node["components"].append({"id": cid_map[c], "label": comps[c]["label"], "icon": comps[c]["icon"]})
            return node
        top_blocks = [build_block(bid) for bid, p in bparent.items() if p is None]
        loose = [{"id": cid_map[c], "label": comps[c]["label"], "icon": comps[c]["icon"]}
                 for c, p in cparent.items() if p is None]
        # connections: remap edge endpoints to nearest component
        centers = {cid_map[cid]: (c["rect"][0]+c["rect"][2]/2, c["rect"][1]+c["rect"][3]/2)
                   for cid, c in comps.items()}
        def resolve(eid):
            if eid in comps: return cid_map[eid]
            if eid in geo:
                r = absr(eid); ec = (r[0]+r[2]/2, r[1]+r[3]/2); best, bd = None, 1e18
                for k, cc in centers.items():
                    d = math.hypot(ec[0]-cc[0], ec[1]-cc[1])
                    if d < bd: bd, best = d, k
                return best if bd < 350 else None
            return None
        conns, seen = [], set()
        for c in cells.values():
            if c.get("edge") != "1": continue
            a, b = resolve(c.get("source")), resolve(c.get("target"))
            if not a or not b or a == b: continue
            key = tuple(sorted((a, b)))
            if key in seen: continue
            seen.add(key); e = {"source": a, "target": b}
            lab = _label(c.get("value"))
            if lab and not junk(lab): e["label"] = lab
            conns.append(e)
        # prune empty + collapse single-child chains + MERGE same-label nesting
        # (e.g. Subaccount-inside-Subaccount -> one Subaccount)
        def collapse(b):
            b["blocks"] = [collapse(s) for s in b["blocks"]]
            b["blocks"] = [s for s in b["blocks"] if s["components"] or s["blocks"]]
            # absorb a child block that repeats this block's label
            merged = []
            for s in b["blocks"]:
                if s["label"] and s["label"] == b["label"]:
                    b["components"] += s["components"]; merged += s["blocks"]
                else:
                    merged.append(s)
            b["blocks"] = merged
            while not b["components"] and len(b["blocks"]) == 1:
                child = b["blocks"][0]
                label = b["label"] if b["label"] else child["label"]
                b = child; b["label"] = label
            return b
        top_blocks = [collapse(b) for b in top_blocks]
        top_blocks = [b for b in top_blocks if b["components"] or b["blocks"]]
        if loose:
            top_blocks.insert(0, {"id": "bext", "label": "External", "blocks": [], "components": loose})
        out.append({"title": title or pname, "blocks": top_blocks, "connections": conns})
    return out

if __name__ == "__main__":
    import sys
    for s in extract_hier(sys.argv[1], title=sys.argv[2] if len(sys.argv)>2 else None):
        print(json.dumps(s, indent=2, ensure_ascii=False))
        def cnt(bs):
            n=len(bs); c=sum(len(b["components"]) for b in bs)
            for b in bs: sub=cnt(b["blocks"]); n+=sub[0]; c+=sub[1]
            return n,c
        nb,nc=cnt(s["blocks"]); print(f"\n# blocks={nb} components={nc} connections={len(s['connections'])}")