|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +add-shape.py — Add a new vertex shape to an existing .drawio diagram file. |
| 4 | +
|
| 5 | +Usage: |
| 6 | + python scripts/add-shape.py <diagram.drawio> <label> <x> <y> [options] |
| 7 | +
|
| 8 | +Examples: |
| 9 | + python scripts/add-shape.py docs/flowchart.drawio "New Step" 400 300 |
| 10 | + python scripts/add-shape.py docs/arch.drawio "Decision" 400 400 \\ |
| 11 | + --width 160 --height 80 \\ |
| 12 | + --style "rhombus;whiteSpace=wrap;html=1;fillColor=#fff2cc;strokeColor=#d6b656;" |
| 13 | + python scripts/add-shape.py docs/arch.drawio "Preview Node" 200 200 --dry-run |
| 14 | +""" |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import argparse |
| 18 | +import hashlib |
| 19 | +import sys |
| 20 | +import time |
| 21 | +import xml.etree.ElementTree as ET |
| 22 | +from pathlib import Path |
| 23 | + |
| 24 | + |
| 25 | +DEFAULT_STYLE = "rounded=1;whiteSpace=wrap;html=1;" |
| 26 | + |
| 27 | + |
| 28 | +def _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 | + |
| 48 | +def _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 | + |
| 54 | +def 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 | + |
| 158 | +def _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 | + |
| 188 | +def 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 | + |
| 212 | +if __name__ == "__main__": |
| 213 | + sys.exit(main()) |
0 commit comments