Maintain groups when importing svg files

manim/mobject/svg/svg_mobject.py

@@ -21,7 +21,7 @@
 from ..geometry.line import Line
 from ..geometry.polygram import Polygon, Rectangle, RoundedRectangle
 from ..opengl.opengl_compatibility import ConvertToOpenGL
-from ..types.vectorized_mobject import VMobject
+from ..types.vectorized_mobject import VGroup, VMobject
 
 __all__ = ["SVGMobject", "VMobjectFromSVGPath"]
 
@@ -127,6 +127,7 @@
         self.stroke_color = stroke_color
         self.stroke_opacity = stroke_opacity  # type: ignore[assignment]
         self.stroke_width = stroke_width  # type: ignore[assignment]
+        self.id_to_vgroup_dict: dict[str, VGroup] = {}
         if self.stroke_width is None:
             self.stroke_width = 0
 
@@ -203,8 +204,9 @@
         svg = se.SVG.parse(modified_file_path)
         modified_file_path.unlink()
 
-        mobjects = self.get_mobjects_from(svg)
+        mobjects, mobject_dict = self.get_mobjects_from(svg)
         self.add(*mobjects)
+        self.id_to_vgroup_dict = mobject_dict
         self.flip(RIGHT)  # Flip y
 
     def get_file_path(self) -> Path:
@@ -258,7 +260,9 @@
                 result[svg_key] = str(svg_default_dict[style_key])
         return result
 
-    def get_mobjects_from(self, svg: se.SVG) -> list[VMobject]:
+    def get_mobjects_from(
+        self, svg: se.SVG
+    ) -> tuple[list[VMobject], dict[str, VGroup]]:
         """Convert the elements of the SVG to a list of mobjects.
 
         Parameters
@@ -267,36 +271,79 @@
             The parsed SVG file.
         """
         result: list[VMobject] = []
-        for shape in svg.elements():
-            # can we combine the two continue cases into one?
-            if isinstance(shape, se.Group):  # noqa: SIM114
-                continue
-            elif isinstance(shape, se.Path):
-                mob: VMobject = self.path_to_mobject(shape)
-            elif isinstance(shape, se.SimpleLine):
-                mob = self.line_to_mobject(shape)
-            elif isinstance(shape, se.Rect):
-                mob = self.rect_to_mobject(shape)
-            elif isinstance(shape, (se.Circle, se.Ellipse)):
-                mob = self.ellipse_to_mobject(shape)
-            elif isinstance(shape, se.Polygon):
-                mob = self.polygon_to_mobject(shape)
-            elif isinstance(shape, se.Polyline):
-                mob = self.polyline_to_mobject(shape)
-            elif isinstance(shape, se.Text):
-                mob = self.text_to_mobject(shape)
-            elif isinstance(shape, se.Use) or type(shape) is se.SVGElement:
-                continue
-            else:
-                logger.warning(f"Unsupported element type: {type(shape)}")
-                continue
-            if mob is None or not mob.has_points():
-                continue
-            self.apply_style_to_mobject(mob, shape)
-            if isinstance(shape, se.Transformable) and shape.apply:
-                self.handle_transform(mob, shape.transform)
-            result.append(mob)
-        return result
+        stack: list[tuple[se.SVGElement, int]] = []
+        stack.append((svg, 1))
+        group_id_number = 0
+        vgroup_stack: list[str] = ["root"]
+        vgroup_names: list[str] = ["root"]
+        vgroups: dict[str, VGroup] = {"root": VGroup()}
+        while len(stack) > 0:
+            element, depth = stack.pop()
+            # Reduce stack heights
+            vgroup_stack = vgroup_stack[0:(depth)]
+            try:
+                group_name = str(element.values["id"])
+            except Exception:
+                group_name = f"numbered_group_{group_id_number}"
+                group_id_number += 1
+            vg = VGroup()
+            vgroup_names.append(group_name)
+            vgroup_stack.append(group_name)
+            vgroups[group_name] = vg
+
+            if isinstance(element, (se.Group, se.Use)):
+                for subelement in element[::-1]:
+                    stack.append((subelement, depth + 1))
+            # Add element to the parent vgroup
+            try:
+                if isinstance(
+                    element,
+                    (
+                        se.Path,
+                        se.SimpleLine,
+                        se.Rect,
+                        se.Circle,
+                        se.Ellipse,
+                        se.Polygon,
+                        se.Polyline,
+                        se.Text,
+                    ),
+                ):
+                    mob = self.get_mob_from_shape_element(element)
+                    if mob is not None:
+                        result.append(mob)
+                        parent_name = vgroup_stack[-2]
+                        for parent_name in vgroup_stack[:-2]:
+                            vgroups[parent_name].add(mob)
+            except Exception as e:
+                print(e)
+
+        return result, vgroups
+
+    def get_mob_from_shape_element(self, shape: se.SVGElement) -> VMobject | None:
+        if isinstance(shape, se.Path):
+            mob: VMobject | None = self.path_to_mobject(shape)
+        elif isinstance(shape, se.SimpleLine):
+            mob = self.line_to_mobject(shape)
+        elif isinstance(shape, se.Rect):
+            mob = self.rect_to_mobject(shape)
+        elif isinstance(shape, (se.Circle, se.Ellipse)):
+            mob = self.ellipse_to_mobject(shape)
+        elif isinstance(shape, se.Polygon):
+            mob = self.polygon_to_mobject(shape)
+        elif isinstance(shape, se.Polyline):
+            mob = self.polyline_to_mobject(shape)
+        elif isinstance(shape, se.Text):
+            mob = self.text_to_mobject(shape)
+        else:
+            logger.warning(f"Unsupported element type: {type(shape)}")
+            mob = None
+        if mob is None or not mob.has_points():
+            return mob
+        self.apply_style_to_mobject(mob, shape)
+        if isinstance(shape, se.Transformable) and shape.apply:
+            self.handle_transform(mob, shape.transform)
+        return mob
 
     @staticmethod
     def handle_transform(mob: VMobject, matrix: se.Matrix) -> VMobject: