Setting the file. One moment. scripts/add-shape.py
Python·213 lines·6 KB
from
__future__
import
annotations
16
17import argparse
18import hashlib
19import sys
20import time
21import xml.etree.ElementTree as ET
22from pathlib import Path
23
24
25DEFAULT_STYLE = "rounded=1;whiteSpace=wrap;html=1;"
26
27
28def _indent_xml(elem: ET.Element, level: int = 0) -> None:
29 """Indent XML tree in-place. Replaces ET.indent() for Python 3.8 compatibility."""
30 indent = "\n" + " " * level
31 if len(elem):
32 if not elem.text or not elem.text.strip():
33 elem.text = indent + " "
34 if not elem.tail or not elem.tail.strip():
35 elem.tail = indent
36 for child in elem:
37 _indent_xml(child, level + 1)
38 # last child tail
39 if not child.tail or not child.tail.strip():
40 child.tail = indent
41 else:
42 if level and (not elem.tail or not elem.tail.strip()):
43 elem.tail = indent
44 if not level:
45 elem.tail = "\n"
46
47
48def _generate_id(label: str, x: int, y: int) -> str:
49 """Generate a short deterministic-ish id based on label + position + time."""
50 seed = f"{label}:{x}:{y}:{time.time_ns()}"
51 return "auto_" + hashlib.sha1(seed.encode()).hexdigest()[:8]
52
53
54def add_shape(
55 path: Path,
56 label: str,
57 x: int,
58 y: int,
59 width: int = 120,
60 height: int = 60,
61 style: str = DEFAULT_STYLE,
62 diagram_index: int = 0,
63 dry_run: bool = False,
64) -> int:
65 """
66 Parse the .drawio file, insert a new vertex cell into the specified diagram page,
67 and write the file back (unless dry_run is True).
68
69 Returns:
70 0 on success, 1 on failure.
71 """
72 # Preserve the original XML declaration / indentation by writing raw bytes.
73 ET.register_namespace("", "")
74
75 try:
76 tree = ET.parse(path)
77 except ET.ParseError as exc:
78 print(f"ERROR: XML parse error in '{path}': {exc}")
79 return 1
80
81 mxfile = tree.getroot()
82 if mxfile.tag != "mxfile":
83 print(f"ERROR: Root element must be <mxfile>, got <{mxfile.tag}>")
84 return 1
85
86 diagrams = mxfile.findall("diagram")
87 if diagram_index >= len(diagrams):
88 print(
89 f"ERROR: diagram-index {diagram_index} is out of range "
90 f"(file has {len(diagrams)} diagram(s))"
91 )
92 return 1
93
94 diagram = diagrams[diagram_index]
95 graph_model = diagram.find("mxGraphModel")
96 if graph_model is None:
97 print(
98 "ERROR: <mxGraphModel> not found as direct child. "
99 "Compressed diagrams are not supported."
100 )
101 return 1
102
103 root_elem = graph_model.find("root")
104 if root_elem is None:
105 print("ERROR: <root> element not found inside <mxGraphModel>")
106 return 1
107
108 # Determine parent id — default to "1" (the default layer)
109 parent_id = "1"
110 existing_ids = {c.get("id") for c in root_elem.findall("mxCell") if c.get("id")}
111 if parent_id not in existing_ids:
112 # Fallback to the first cell id that isn't "0"
113 for c in root_elem.findall("mxCell"):
114 cid = c.get("id")
115 if cid and cid != "0":
116 parent_id = cid
117 break
118
119 # Generate a unique id
120 new_id = _generate_id(label, x, y)
121 while new_id in existing_ids:
122 new_id = _generate_id(label + "_", x, y)
123
124 # Build the new mxCell element
125 new_cell = ET.Element("mxCell")
126 new_cell.set("id", new_id)
127 new_cell.set("value", label)
128 new_cell.set("style", style)
129 new_cell.set("vertex", "1")
130 new_cell.set("parent", parent_id)
131
132 geom = ET.SubElement(new_cell, "mxGeometry")
133 geom.set("x", str(x))
134 geom.set("y", str(y))
135 geom.set("width", str(width))
136 geom.set("height", str(height))
137 geom.set("as", "geometry")
138
139 if dry_run:
140 print("DRY RUN — new cell XML (not written):")
141 print(ET.tostring(new_cell, encoding="unicode"))
142 print(f"\nWould add to diagram '{diagram.get('name', diagram_index)}' in '{path}'")
143 return 0
144
145 root_elem.append(new_cell)
146
147 # Write back preserving XML declaration (uses _indent_xml for Python 3.8 compat)
148 _indent_xml(tree.getroot())
149 tree.write(str(path), encoding="utf-8", xml_declaration=True)
150
151 print(
152 f"Added shape id=\"{new_id}\" to page {diagram_index} "
153 f"('{diagram.get('name', '')}') of {path}"
154 )
155 return 0
156
157
158def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
159 parser = argparse.ArgumentParser(
160 description="Add a shape to an existing .drawio diagram file.",
161 formatter_class=argparse.RawDescriptionHelpFormatter,
162 )
163 parser.add_argument("diagram", help="Path to the .drawio file")
164 parser.add_argument("label", help="Text label for the new shape")
165 parser.add_argument("x", type=int, help="X coordinate (pixels)")
166 parser.add_argument("y", type=int, help="Y coordinate (pixels)")
167 parser.add_argument("--width", type=int, default=120, help="Shape width (default: 120)")
168 parser.add_argument("--height", type=int, default=60, help="Shape height (default: 60)")
169 parser.add_argument(
170 "--style",
171 default=DEFAULT_STYLE,
172 help=f'draw.io style string (default: "{DEFAULT_STYLE}")',
173 )
174 parser.add_argument(
175 "--diagram-index",
176 type=int,
177 default=0,
178 help="0-based index of the diagram page to add to (default: 0)",
179 )
180 parser.add_argument(
181 "--dry-run",
182 action="store_true",
183 help="Print the new cell XML without writing to file",
184 )
185 return parser.parse_args(argv)
186
187
188def main(argv: list[str] | None = None) -> int:
189 args = _parse_args(argv)
190 path = Path(args.diagram)
191
192 if not path.exists():
193 print(f"ERROR: File not found: {path}")
194 return 1
195 if not path.is_file():
196 print(f"ERROR: Not a file: {path}")
197 return 1
198
199 return add_shape(
200 path=path,
201 label=args.label,
202 x=args.x,
203 y=args.y,
204 width=args.width,
205 height=args.height,
206 style=args.style,
207 diagram_index=args.diagram_index,
208 dry_run=args.dry_run,
209 )
210
211
212if __name__ == "__main__":
213 sys.exit(main())