Add another simplification and reuse mapping for geometries

Author: vanous

pygdtf/__init__.py

@@ -268,78 +268,12 @@ def _read_xml(self):
         self.geometries = Geometries()
         geometry_collect = self._root.find("Geometries")
         if geometry_collect is not None:
-            for i in geometry_collect.findall("Geometry"):
-                self.geometries.append(
-                    Geometry(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Axis"):
-                self.geometries.append(
-                    GeometryAxis(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("FilterBeam"):
-                self.geometries.append(
-                    GeometryFilterBeam(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("FilterColor"):
-                self.geometries.append(
-                    GeometryFilterColor(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("FilterGobo"):
-                self.geometries.append(
-                    GeometryFilterGobo(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("FilterShaper"):
-                self.geometries.append(
-                    GeometryFilterShaper(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("MediaServerMaster"):
-                self.geometries.append(
-                    GeometryMediaServerMaster(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("MediaServerLayer"):
-                self.geometries.append(
-                    GeometryMediaServerLayer(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("MediaServerCamera"):
-                self.geometries.append(
-                    GeometryMediaServerCamera(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Inventory"):
-                self.geometries.append(
-                    GeometryInventory(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Beam"):
-                self.geometries.append(
-                    GeometryBeam(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("WiringObject"):
-                self.geometries.append(
-                    GeometryWiringObject(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("GeometryReference"):
-                self.geometries.append(
-                    GeometryReference(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Laser"):
-                self.geometries.append(
-                    GeometryLaser(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Support"):
-                self.geometries.append(
-                    GeometrySupport(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Structure"):
-                self.geometries.append(
-                    GeometryStructure(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Display"):
-                self.geometries.append(
-                    GeometryDisplay(xml_node=i, xml_parent=geometry_collect)
-                )
-            for i in geometry_collect.findall("Magnet"):
-                self.geometries.append(
-                    GeometryMagnet(xml_node=i, xml_parent=geometry_collect)
-                )
+            for geometry_node in list(geometry_collect):
+                cls = TAG_TO_GEOMETRY_CLASS.get(geometry_node.tag)
+                if cls:
+                    self.geometries.append(
+                        cls(xml_node=geometry_node, xml_parent=geometry_collect)
+                    )
 
         dmx_mode_collect = self._root.find("DMXModes")
         if dmx_mode_collect is not None:

pygdtf/geometries.py

@@ -113,40 +113,21 @@ def _read_xml(self, xml_node: "Element", xml_parent: Optional["Element"] = None)
 
         self.model = xml_node.attrib.get("Model")
         self.position = Matrix(xml_node.attrib.get("Position", 0))
-        tag_to_class = {
-            "Geometry": Geometry,
-            "Axis": GeometryAxis,
-            "FilterBeam": GeometryFilterBeam,
-            "FilterColor": GeometryFilterColor,
-            "FilterGobo": GeometryFilterGobo,
-            "FilterShaper": GeometryFilterShaper,
-            "MediaServerMaster": GeometryMediaServerMaster,
-            "MediaServerLayer": GeometryMediaServerLayer,
-            "MediaServerCamera": GeometryMediaServerCamera,
-            "Inventory": GeometryInventory,
-            "Beam": GeometryBeam,
-            "WiringObject": GeometryWiringObject,
-            "GeometryReference": GeometryReference,
-            "Laser": GeometryLaser,
-            "Structure": GeometryStructure,
-            "Support": GeometrySupport,
-            "Magnet": GeometryMagnet,
-            "Display": GeometryDisplay,
-        }
         for child in list(xml_node):
-            cls = tag_to_class.get(child.tag)
+            cls = TAG_TO_GEOMETRY_CLASS.get(child.tag)
             if cls:
                 self.geometries.append(cls(xml_node=child, xml_parent=xml_node))
 
     def __str__(self):
         return f"{self.name} ({self.model})"
 
     def _tag(self):
-        if type(self) is Geometry:
-            return "Geometry"
-        if self.__class__.__name__.startswith("Geometry"):
-            return self.__class__.__name__.replace("Geometry", "", 1)
-        return type(self).__name__
+        tag = CLASS_TO_GEOMETRY_TAG.get(type(self))
+        if tag:
+            return tag
+        raise KeyError(
+            f"Geometry class not registered in mapping: {type(self).__name__}"
+        )
 
     def to_xml(self):
         attrs = {}
@@ -572,3 +553,26 @@ def _read_xml(self, xml_node: "Element", xml_parent: Optional["Element"] = None)
 
     def __str__(self):
         return f"{self.name} ({self.model})"
+
+
+TAG_TO_GEOMETRY_CLASS = {
+    "Geometry": Geometry,
+    "Axis": GeometryAxis,
+    "FilterBeam": GeometryFilterBeam,
+    "FilterColor": GeometryFilterColor,
+    "FilterGobo": GeometryFilterGobo,
+    "FilterShaper": GeometryFilterShaper,
+    "MediaServerMaster": GeometryMediaServerMaster,
+    "MediaServerLayer": GeometryMediaServerLayer,
+    "MediaServerCamera": GeometryMediaServerCamera,
+    "Inventory": GeometryInventory,
+    "Beam": GeometryBeam,
+    "WiringObject": GeometryWiringObject,
+    "GeometryReference": GeometryReference,
+    "Laser": GeometryLaser,
+    "Structure": GeometryStructure,
+    "Support": GeometrySupport,
+    "Magnet": GeometryMagnet,
+    "Display": GeometryDisplay,
+}
+CLASS_TO_GEOMETRY_TAG = {cls: tag for tag, cls in TAG_TO_GEOMETRY_CLASS.items()}