diff --git a/etc/development scripts/import_export.py b/etc/development scripts/import_export.py new file mode 100644 index 0000000..82fd676 --- /dev/null +++ b/etc/development scripts/import_export.py @@ -0,0 +1,21 @@ +from pathlib import Path + +import drawpyo +from drawpyo import load_diagram + +# Load Draw.io diagram +relative_path = Path("..") / "reference drawio charts" / "object.drawio" +file_path = (Path(__file__).parent / relative_path).resolve() + +diagram = load_diagram(file_path) + +# Create file & page +file = drawpyo.File() +file.file_path = str(Path.home() / "Test Drawpyo Charts") +file.file_name = "Test_object.drawio" + +page = drawpyo.Page(file=file) + +diagram.add_to(page) +# Write the file +file.write() diff --git a/etc/reference drawio charts/object.drawio b/etc/reference drawio charts/object.drawio new file mode 100644 index 0000000..ea38dd4 --- /dev/null +++ b/etc/reference drawio charts/object.drawio @@ -0,0 +1,47 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/drawpyo/diagram/edges.py b/src/drawpyo/diagram/edges.py index ca5d8a1..85b6a6f 100644 --- a/src/drawpyo/diagram/edges.py +++ b/src/drawpyo/diagram/edges.py @@ -92,6 +92,13 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.xml_class: str = "mxCell" + self.object_attributes: Dict[str, str] = dict( + kwargs.get("object_attributes", {}) + ) + self.user_object_attributes: Dict[str, str] = dict( + kwargs.get("user_object_attributes", {}) + ) + # Style self.color_scheme: Optional[ColorScheme] = kwargs.get("color_scheme", None) self.text_format: Optional[TextFormat] = kwargs.get("text_format", TextFormat()) @@ -210,15 +217,20 @@ def attributes(self) -> Dict[str, Any]: Returns: dict: Dictionary of object attributes and their values """ + id_value = self.id + if self.object_attributes or self.user_object_attributes: + id_value = None base_attr_dict: Dict[str, Any] = { - "id": self.id, + "id": id_value, "style": self.style, "edge": self.edge, "parent": self.xml_parent_id, "source": self.source_id, "target": self.target_id, } - if self.value is not None: + if self.value is not None and not ( + self.object_attributes or self.user_object_attributes + ): base_attr_dict["value"] = self.value return base_attr_dict @@ -591,8 +603,36 @@ def xml(self) -> str: tag: str = ( self.xml_open_tag + "\n " + self.geometry.xml + "\n" + self.xml_close_tag ) + wrapper_tag = None + wrapper_attrs = None + if self.user_object_attributes: + wrapper_tag = "UserObject" + wrapper_attrs = self.user_object_attributes + elif self.object_attributes: + wrapper_tag = "object" + wrapper_attrs = self.object_attributes + if wrapper_tag and wrapper_attrs is not None: + return ( + self._wrapper_open_tag(wrapper_tag, wrapper_attrs) + + "\n " + + tag.replace("\n", "\n ") + + f"\n" + ) return tag + def _wrapper_open_tag(self, tag: str, attrs: Dict[str, str]) -> str: + attrs = {k: v for k, v in attrs.items() if v is not None} + if "label" not in attrs and self.value is not None: + attrs["label"] = self.value + if "id" not in attrs: + attrs["id"] = self.id + + open_tag = f"<{tag}" + for att, value in attrs.items(): + xml_parameter = self.xml_ify(str(value)) + open_tag = open_tag + " " + att + '="' + xml_parameter + '"' + return open_tag + ">" + class BasicEdge(Edge): pass diff --git a/src/drawpyo/diagram/objects.py b/src/drawpyo/diagram/objects.py index 20dfb59..8060059 100644 --- a/src/drawpyo/diagram/objects.py +++ b/src/drawpyo/diagram/objects.py @@ -1,15 +1,15 @@ from os import path -from typing import Optional, Dict, Any, List, Union, Tuple -from ..utils.logger import logger +from typing import Any, Dict, List, Optional, Tuple, Union +from ..utils.color_scheme import ColorScheme +from ..utils.logger import logger +from ..utils.standard_colors import StandardColor from .base_diagram import ( DiagramBase, Geometry, import_shape_database, ) from .text_format import TextFormat -from ..utils.color_scheme import ColorScheme -from ..utils.standard_colors import StandardColor __all__ = ["Object", "BasicObject", "Group", "object_from_library"] @@ -143,6 +143,12 @@ def __init__( self.autosize_to_children: bool = kwargs.get("autosize_to_children", False) self.autocontract: bool = kwargs.get("autocontract", False) self.autosize_margin: int = kwargs.get("autosize_margin", 20) + self.object_attributes: Dict[str, str] = dict( + kwargs.get("object_attributes", {}) + ) + self.user_object_attributes: Dict[str, str] = dict( + kwargs.get("user_object_attributes", {}) + ) # Geometry self.position: Optional[tuple] = position @@ -352,9 +358,16 @@ def format_as_library_object( @property def attributes(self) -> Dict[str, Any]: + id_value = self.id + value_value = self.value + if not (self.tag or self.tooltip) and ( + self.object_attributes or self.user_object_attributes + ): + id_value = None + value_value = None return { - "id": self.id, - "value": self.value, + "id": id_value, + "value": value_value, "style": self.style, "vertex": self.vertex, "parent": self.xml_parent_id, @@ -704,8 +717,38 @@ def xml(self) -> str: tag: str = ( self.xml_open_tag + "\n " + self.geometry.xml + "\n" + self.xml_close_tag ) + if self.tag or self.tooltip: + return tag + wrapper_tag = None + wrapper_attrs = None + if self.user_object_attributes: + wrapper_tag = "UserObject" + wrapper_attrs = self.user_object_attributes + elif self.object_attributes: + wrapper_tag = "object" + wrapper_attrs = self.object_attributes + if wrapper_tag and wrapper_attrs is not None: + return ( + self._wrapper_open_tag(wrapper_tag, wrapper_attrs) + + "\n " + + tag.replace("\n", "\n ") + + f"\n" + ) return tag + def _wrapper_open_tag(self, tag: str, attrs: Dict[str, str]) -> str: + attrs = {k: v for k, v in attrs.items() if v is not None} + if "label" not in attrs and self.value is not None: + attrs["label"] = self.value + if "id" not in attrs: + attrs["id"] = self.id + + open_tag = f"<{tag}" + for att, value in attrs.items(): + xml_parameter = self.xml_ify(str(value)) + open_tag = open_tag + " " + att + '="' + xml_parameter + '"' + return open_tag + ">" + class BasicObject(Object): pass diff --git a/src/drawpyo/drawio_import/drawio_parser.py b/src/drawpyo/drawio_import/drawio_parser.py index 3a1b625..94d2a3d 100644 --- a/src/drawpyo/drawio_import/drawio_parser.py +++ b/src/drawpyo/drawio_import/drawio_parser.py @@ -1,11 +1,12 @@ import xml.etree.ElementTree as ET +from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional -from dataclasses import dataclass, field -from .raw import RawMxCell, RawGeometry -from drawpyo import logger -from drawpyo.diagram import Object, Edge, DiagramBase +from drawpyo import Page, logger +from drawpyo.diagram import DiagramBase, Edge, Object + +from .raw import RawGeometry, RawMxCell # ----------------------------- @@ -35,10 +36,117 @@ def element_count(self) -> int: """Total number of elements (shapes + edges).""" return len(self.shapes) + len(self.edges) + def add_to( + self, + page: Page, + offset: tuple = (0, 0), + include_shapes: bool = True, + include_edges: bool = True, + ) -> None: + """Attach diagram elements to a page, optionally applying an offset. + + Args: + page: Target Page instance to attach elements to + offset: (dx, dy) applied to top-level objects and edge points + include_shapes: Whether to add shapes to the page + include_edges: Whether to add edges to the page + """ + dx, dy = offset + + if include_shapes: + for shape in self.shapes: + _update_page_links(shape, page) + shape.page = page + if dx or dy: + for shape in self.shapes: + if isinstance(shape, Object) and shape.parent is None: + shape.position = ( + shape.position[0] + dx, + shape.position[1] + dy, + ) + + if include_edges: + for edge in self.edges: + _update_page_links(edge, page) + edge.page = page + if dx or dy: + for edge in self.edges: + for point in edge.geometry.points: + point.x = point.x + dx + point.y = point.y + dy + # ----------------------------- # XML Parsing # ----------------------------- +def _build_raw_cell( + cell_elem: ET.Element, cell_id_override: Optional[str] = None +) -> Optional[RawMxCell]: + """ + Builds a RawMxCell object from an XML element. + + Args: + cell_elem: The XML element representing a cell. + cell_id_override: An optional override for the cell ID. + + Returns: + A RawMxCell object or None if the cell ID is missing. + """ + cell_id = cell_id_override or cell_elem.get("id") + if not cell_id: + return None + + cell = RawMxCell( + id=cell_id, + parent=cell_elem.get("parent"), + value=cell_elem.get("value"), + style=cell_elem.get("style"), + is_vertex=cell_elem.get("vertex") == "1", + is_edge=cell_elem.get("edge") == "1", + source=cell_elem.get("source"), + target=cell_elem.get("target"), + ) + + geo_elem = cell_elem.find("mxGeometry") + if geo_elem is not None: + points = [] + + points_array = geo_elem.find("Array[@as='points']") + if points_array is not None: + for point_elem in points_array.findall("mxPoint"): + x = point_elem.get("x") + y = point_elem.get("y") + if x is not None and y is not None: + points.append((float(x), float(y))) + + cell.geometry = RawGeometry( + x=float(geo_elem.get("x")) if geo_elem.get("x") else None, + y=float(geo_elem.get("y")) if geo_elem.get("y") else None, + width=float(geo_elem.get("width")) if geo_elem.get("width") else None, + height=(float(geo_elem.get("height")) if geo_elem.get("height") else None), + relative=geo_elem.get("relative") == "1", + points=points, + ) + + return cell + + +def _update_page_links(element: DiagramBase, page: Page) -> None: + """ + Updates the 'link' attribute of user object attributes if it points to a Draw.io page. + + Args: + element: The diagram element (shape or edge) to update. + page: The target Page instance. + """ + user_attrs = getattr(element, "user_object_attributes", None) + if not user_attrs: + return + link = user_attrs.get("link") + if link and link.startswith("data:page/id,"): + user_attrs["link"] = f"data:page/id,{page.diagram.id}" + + def _parse_drawio_xml(xml_string: str) -> Dict[str, RawMxCell]: """Parses draw.io XML into a dictionary of RawMxCell objects keyed by their IDs. @@ -54,44 +162,29 @@ def _parse_drawio_xml(xml_string: str) -> Dict[str, RawMxCell]: root = ET.fromstring(xml_string) cells: Dict[str, RawMxCell] = {} - for cell_elem in root.findall(".//mxCell"): - cell_id = cell_elem.get("id") - if not cell_id: - continue + for wrapper_tag in ("object", "UserObject"): + for wrapper_elem in root.findall(f".//{wrapper_tag}"): + cell_elem = wrapper_elem.find("mxCell") + if cell_elem is None: + continue - cell = RawMxCell( - id=cell_id, - parent=cell_elem.get("parent"), - value=cell_elem.get("value"), - style=cell_elem.get("style"), - is_vertex=cell_elem.get("vertex") == "1", - is_edge=cell_elem.get("edge") == "1", - source=cell_elem.get("source"), - target=cell_elem.get("target"), - ) + wrapper_attrs = dict(wrapper_elem.attrib) + cell_id = cell_elem.get("id") or wrapper_attrs.get("id") + cell = _build_raw_cell(cell_elem, cell_id_override=cell_id) + if cell is None: + continue + + if wrapper_tag == "UserObject": + cell.user_object_attributes = wrapper_attrs + else: + cell.object_attributes = wrapper_attrs - geo_elem = cell_elem.find("mxGeometry") - if geo_elem is not None: - points = [] - - points_array = geo_elem.find("Array[@as='points']") - if points_array is not None: - for point_elem in points_array.findall("mxPoint"): - x = point_elem.get("x") - y = point_elem.get("y") - if x is not None and y is not None: - points.append((float(x), float(y))) - - cell.geometry = RawGeometry( - x=float(geo_elem.get("x")) if geo_elem.get("x") else None, - y=float(geo_elem.get("y")) if geo_elem.get("y") else None, - width=float(geo_elem.get("width")) if geo_elem.get("width") else None, - height=( - float(geo_elem.get("height")) if geo_elem.get("height") else None - ), - relative=geo_elem.get("relative") == "1", - points=points, - ) + cells[cell.id] = cell + + for cell_elem in root.findall(".//mxCell"): + cell = _build_raw_cell(cell_elem) + if cell is None or cell.id in cells: + continue cells[cell.id] = cell @@ -142,7 +235,11 @@ def _build_vertices(raw_cells: Dict[str, RawMxCell]) -> Dict[str, DiagramBase]: if not cell.is_vertex: continue - obj = Object(value=cell.value) + obj = Object( + id=cell.id, + object_attributes=cell.object_attributes, + user_object_attributes=cell.user_object_attributes, + ) if cell.style: obj.apply_style_string(cell.style) @@ -169,7 +266,13 @@ def _apply_geometry_recursive( raw_cells: Dict[str, RawMxCell], elements: Dict[str, DiagramBase], ): - """Apply geometry recursively, preserving relative positions.""" + """Apply geometry recursively, preserving relative positions. + + Args: + cell_id: The ID of the current cell to process. + raw_cells: Dictionary of all raw cell data. + elements: Dictionary of DiagramBase objects, to which geometry will be applied. + """ cell = raw_cells[cell_id] obj = elements[cell_id] @@ -198,7 +301,11 @@ def _build_edges(raw_cells: Dict[str, RawMxCell], elements: Dict[str, DiagramBas if not cell.is_edge: continue - e = Edge() + e = Edge( + id=cell.id, + object_attributes=cell.object_attributes, + user_object_attributes=cell.user_object_attributes, + ) if cell.style: e.apply_style_string(cell.style) diff --git a/src/drawpyo/drawio_import/raw.py b/src/drawpyo/drawio_import/raw.py index 0cbc843..e6fbdef 100644 --- a/src/drawpyo/drawio_import/raw.py +++ b/src/drawpyo/drawio_import/raw.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Optional, List, Tuple +from typing import Optional, List, Tuple, Dict @dataclass @@ -21,6 +21,9 @@ class RawMxCell: value: Optional[str] = None style: Optional[str] = None + object_attributes: Dict[str, str] = field(default_factory=dict) + user_object_attributes: Dict[str, str] = field(default_factory=dict) + is_vertex: bool = False is_edge: bool = False diff --git a/tests/import_tests/test_import_export_object_userobject.py b/tests/import_tests/test_import_export_object_userobject.py new file mode 100644 index 0000000..98f028b --- /dev/null +++ b/tests/import_tests/test_import_export_object_userobject.py @@ -0,0 +1,121 @@ +import re + +import pytest + +from drawpyo import File, Page, load_diagram +from drawpyo.diagram import Edge, Object + +SAMPLE_XML = """ + + + + + + + + + + + + + + + + + + + + + + + + + + +""" + + +@pytest.fixture +def diagram(tmp_path): + file_path = tmp_path / "object_userobject.drawio" + file_path.write_text(SAMPLE_XML) + return load_diagram(str(file_path)) + + +def test_object_userobject_import(diagram): + assert diagram is not None + assert len(diagram.shapes) == 2 + assert len(diagram.edges) == 2 + + wrapped_edge = diagram.get_by_id("2xg5nJA0Zl71o69cM8dW-4") + wrapped_obj = diagram.get_by_id("2xg5nJA0Zl71o69cM8dW-1") + user_obj = diagram.get_by_id("2xg5nJA0Zl71o69cM8dW-2") + + assert isinstance(wrapped_edge, Edge) + assert isinstance(wrapped_obj, Object) + assert isinstance(user_obj, Object) + + assert wrapped_edge.object_attributes.get("test") == "test" + assert wrapped_obj.object_attributes.get("test") == "test" + assert ( + user_obj.user_object_attributes.get("link") + == "data:page/id,hyHWzmBswkbCvm2kR1NW" + ) + + +def test_object_userobject_export(diagram): + file = File(file_name="test.drawio", file_path=".") + page = Page(file=file) + + diagram.add_to(page) + + xml = file.xml + + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-4")(?=[^>]*\btest="test")[^>]*>', + xml, + ) + assert re.search( + rf']*\bid="2xg5nJA0Zl71o69cM8dW-2")(?=[^>]*\blink=\"data:page/id,{page.diagram.id}\")[^>]*>', + xml, + ) + assert "data:page/id,hyHWzmBswkbCvm2kR1NW" not in xml + assert re.search( + rf']*\bid="{page.diagram.id}"', + xml, + ) + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-4"[^>]*>\s*]*\bid=)', + xml, + ) + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-1"[^>]*>\s*]*\bid=)', + xml, + ) + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-2"[^>]*>\s*]*\bid=)', + xml, + ) + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-4"[^>]*>\s*]*\bvalue=)', + xml, + ) + assert re.search( + r']*\bid="2xg5nJA0Zl71o69cM8dW-2"[^>]*>\s*]*\bvalue=)', + xml, + ) + + +def test_parsed_diagram_add_to_offset(diagram): + file = File(file_name="test.drawio", file_path=".") + page = Page(file=file) + + diagram.add_to(page, offset=(10, 20)) + + obj_1 = diagram.get_by_id("2xg5nJA0Zl71o69cM8dW-1") + obj_2 = diagram.get_by_id("2xg5nJA0Zl71o69cM8dW-2") + + assert obj_1 is not None + assert obj_2 is not None + assert obj_1.position == (360, 390) + assert obj_2.position == (160, 300)