#!/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 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'])}")