File size: 7,398 Bytes
f5dc454 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | #!/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'])}")
|