Refactorization of the generate_style and write_cell functions

write_cell has a new argument: "root".

drawio2py/generator.py

@@ -26,6 +26,83 @@ def generate_mxfile(drawio: DrawIOFile) -> ET.Element:
 		graph.append(dia)
 	return graph
 
+def generate_style(style: Style) -> str:
+	res = ""
+	for k, v in style.data.items():
+		res += "%s=%s;" % (str(k), str(v))
+	return res
+
+def write_cell(cell, root):
+	attrs = {}
+	if not cell.properties:
+		attrs["id"] = cell.id
+	if cell.value:
+		attrs["value"] = cell.value
+	if cell.style:
+		attrs["style"] = generate_style(cell.style)
+	if cell.parent:
+		attrs["parent"] = cell.parent.id
+	if isinstance(cell, Edge):
+		attrs["edge"] = "1"
+		if cell.source:
+			attrs["source"] = cell.source.id
+		if cell.target:
+			attrs["target"] = cell.target.id
+			
+	if cell.properties:
+		# Wrap in <object> if there are properties
+		properties_with_id = cell.properties.copy()
+		properties_with_id['id'] = cell.id  # <object> gets the ID, not the <mxCell> (WTF drawio!)
+		par = ET.SubElement(root, "object", properties_with_id)
+	else:
+		par = root
+
+	# Create the actual <mxCell>
+	c = ET.SubElement(par, "mxCell", attrs, **cell.attributes)
+
+	def optional_attributes(dict):
+		attrs = {}
+		for key, val in dict.items():
+			if val != None:
+				attrs[key] = str(val)
+		return attrs
+
+	def write_point(parent_xml, point: Point, as_str: Optional[str] = None):
+		attrs = optional_attributes({
+			"x": point.x,
+			"y": point.y,
+			"as": as_str,
+		})
+		return ET.SubElement(parent_xml, "mxPoint", attrs)
+
+	def write_geometry(parent_xml, geometry):
+		return ET.SubElement(parent_xml, "mxGeometry", optional_attributes({
+			"x": geometry.x,
+			"y": geometry.y,
+			"width": geometry.width,
+			"height": geometry.height,
+			"relative": "1" if geometry.relative else None,
+			"as": "geometry",
+		}))
+
+	# Geometry
+	if isinstance(cell, Vertex) or isinstance(cell, Edge):
+		g = write_geometry(c, cell.geometry)
+		if cell.geometry.offset != None:
+			write_point(g, cell.geometry.offset, "offset")
+		if isinstance(cell, Edge):
+			if len(cell.geometry.points) > 0:
+				a = ET.SubElement(g, "Array", {"as": "points"})
+				for p in cell.geometry.points:
+					write_point(a, p)
+			if cell.geometry.source_point is not None:
+				write_point(g, cell.geometry.source_point, "sourcePoint")
+			if cell.geometry.target_point is not None:
+				write_point(g, cell.geometry.target_point, "targetPoint")
+
+	for child in cell.children:
+		write_cell(child, root)
+
 # generates XML tree for a drawio "page" (root element: <diagram>)
 def generate_diagram(page: Page) -> ET.Element:
 	dia = ET.Element("diagram", {
@@ -35,84 +112,7 @@ def generate_diagram(page: Page) -> ET.Element:
 	mxgm = ET.SubElement(dia, "mxGraphModel", page.attributes)
 	root = ET.SubElement(mxgm, "root")
 
-	def generate_style(style: Style) -> str:
-		res = ""
-		for k, v in style.data.items():
-			res += "%s=%s;" % (str(k), str(v))
-		return res
-
-	def write_cell(cell):
-		attrs = {}
-		if not cell.properties:
-			attrs["id"] = cell.id
-		if cell.value:
-			attrs["value"] = cell.value
-		if cell.style:
-			attrs["style"] = generate_style(cell.style)
-		if cell.parent:
-			attrs["parent"] = cell.parent.id
-		if isinstance(cell, Edge):
-			attrs["edge"] = "1"
-			if cell.source:
-				attrs["source"] = cell.source.id
-			if cell.target:
-				attrs["target"] = cell.target.id
-				
-		if cell.properties:
-			# Wrap in <object> if there are properties
-			properties_with_id = cell.properties.copy()
-			properties_with_id['id'] = cell.id  # <object> gets the ID, not the <mxCell> (WTF drawio!)
-			par = ET.SubElement(root, "object", properties_with_id)
-		else:
-			par = root
-
-		# Create the actual <mxCell>
-		c = ET.SubElement(par, "mxCell", attrs, **cell.attributes)
-
-		def optional_attributes(dict):
-			attrs = {}
-			for key, val in dict.items():
-				if val != None:
-					attrs[key] = str(val)
-			return attrs
-
-		def write_point(parent_xml, point: Point, as_str: Optional[str] = None):
-			attrs = optional_attributes({
-				"x": point.x,
-				"y": point.y,
-				"as": as_str,
-			})
-			return ET.SubElement(parent_xml, "mxPoint", attrs)
-
-		def write_geometry(parent_xml, geometry):
-			return ET.SubElement(parent_xml, "mxGeometry", optional_attributes({
-				"x": geometry.x,
-				"y": geometry.y,
-				"width": geometry.width,
-				"height": geometry.height,
-				"relative": "1" if geometry.relative else None,
-				"as": "geometry",
-			}))
-
-		# Geometry
-		if isinstance(cell, Vertex) or isinstance(cell, Edge):
-			g = write_geometry(c, cell.geometry)
-			if cell.geometry.offset != None:
-				write_point(g, cell.geometry.offset, "offset")
-			if isinstance(cell, Edge):
-				if len(cell.geometry.points) > 0:
-					a = ET.SubElement(g, "Array", {"as": "points"})
-					for p in cell.geometry.points:
-						write_point(a, p)
-				if cell.geometry.source_point is not None:
-					write_point(g, cell.geometry.source_point, "sourcePoint")
-				if cell.geometry.target_point is not None:
-					write_point(g, cell.geometry.target_point, "targetPoint")
-
-		for child in cell.children:
-			write_cell(child)
-
-	write_cell(page.root)
+	write_cell(page.root, root)
 	return dia
 
 
@@ -139,12 +139,6 @@ def generate_library(lib : Dict[str, Cell], file_object : str) -> None:
 	keys = lib.keys()
 	nbKeys = len(keys)
 
-	def generate_style(style: Style) -> str:
-		res = ""
-		for k, v in style.data.items():
-			res += "%s=%s;" % (str(k), str(v))
-		return res
-
 	for idx, key in enumerate(keys):
 		
 		fileContent += '  {\n'
@@ -160,79 +154,7 @@ def generate_library(lib : Dict[str, Cell], file_object : str) -> None:
 		empty_root = generate_empty_root_and_layer()
 		empty_root.children[0].children.append(lib.get(key))
 
-
-		def write_cell(cell):
-			attrs = {}
-			if not cell.properties:
-				attrs["id"] = cell.id
-			if cell.value:
-				attrs["value"] = cell.value
-			if cell.style:
-				attrs["style"] = generate_style(cell.style)
-			if cell.parent:
-				attrs["parent"] = cell.parent.id
-			if isinstance(cell, Edge):
-				attrs["edge"] = "1"
-				if cell.source:
-					attrs["source"] = cell.source.id
-				if cell.target:
-					attrs["target"] = cell.target.id
-					
-			if cell.properties:
-				# Wrap in <object> if there are properties
-				properties_with_id = cell.properties.copy()
-				properties_with_id['id'] = cell.id  # <object> gets the ID, not the <mxCell> (WTF drawio!)
-				par = ET.SubElement(root, "object", properties_with_id)
-			else:
-				par = root
-
-			# Create the actual <mxCell>
-			c = ET.SubElement(par, "mxCell", attrs, **cell.attributes)
-
-			def optional_attributes(dict):
-				attrs = {}
-				for key, val in dict.items():
-					if val != None:
-						attrs[key] = str(val)
-				return attrs
-
-			def write_point(parent_xml, point: Point, as_str: Optional[str] = None):
-				attrs = optional_attributes({
-					"x": point.x,
-					"y": point.y,
-					"as": as_str,
-				})
-				return ET.SubElement(parent_xml, "mxPoint", attrs)
-
-			def write_geometry(parent_xml, geometry):
-				return ET.SubElement(parent_xml, "mxGeometry", optional_attributes({
-					"x": geometry.x,
-					"y": geometry.y,
-					"width": geometry.width,
-					"height": geometry.height,
-					"relative": "1" if geometry.relative else None,
-					"as": "geometry",
-				}))
-
-			# Geometry
-			if isinstance(cell, Vertex) or isinstance(cell, Edge):
-				g = write_geometry(c, cell.geometry)
-				if cell.geometry.offset != None:
-					write_point(g, cell.geometry.offset, "offset")
-				if isinstance(cell, Edge):
-					if len(cell.geometry.points) > 0:
-						a = ET.SubElement(g, "Array", {"as": "points"})
-						for p in cell.geometry.points:
-							write_point(a, p)
-					if cell.geometry.source_point is not None:
-						write_point(g, cell.geometry.source_point, "sourcePoint")
-					if cell.geometry.target_point is not None:
-						write_point(g, cell.geometry.target_point, "targetPoint")
-
-			for child in cell.children:
-				write_cell(child)
-
-		write_cell(empty_root)
+		write_cell(empty_root, root)
 		
 		xml = ET.tostring(mxgm, encoding="unicode")
 		xml = xml.replace('<', '&lt;')