Clean up Scene.remove function, delete restructure_list_to_exclude_certain_family_members

manimlib/scene/scene.py

@@ -28,7 +28,6 @@ from manimlib.mobject.types.vectorized_mobject import VMobject
 from manimlib.scene.scene_file_writer import SceneFileWriter
 from manimlib.utils.config_ops import digest_config
 from manimlib.utils.family_ops import extract_mobject_family_members
-from manimlib.utils.family_ops import restructure_list_to_exclude_certain_family_members
 
 from typing import TYPE_CHECKING
 
@@ -302,10 +301,32 @@ class Scene(object):
         ))
         return self
 
-    def remove(self, *mobjects_to_remove: Mobject):
+    def replace(self, mobject: Mobject, *replacements: Mobject):
-        self.mobjects = restructure_list_to_exclude_certain_family_members(
+        if mobject in self.mobjects:
-            self.mobjects, mobjects_to_remove
+            index = self.mobjects.index(mobject)
-        )
+            self.mobjects = [
+                *self.mobjects[:index],
+                *replacements,
+                *self.mobjects[index + 1:]
+            ]
+        return self
+
+    def remove(self, *mobjects: Mobject):
+        """
+        Removes anything in mobjects from scenes mobject list, but in the event that one
+        of the items to be removed is a member of the family of an item in mobject_list,
+        the other family members are added back into the list.
+
+        For example, if the scene includes Group(m1, m2, m3), and we call scene.remove(m1),
+        the desired behavior is for the scene to then include m2 and m3 (ungrouped).
+        """
+        for mob in mobjects:
+            # First restructure list so that parents/grandparents/etc. are replaced
+            # with their children
+            for ancestor in reversed(mob.get_ancestors()):
+                self.replace(ancestor, *ancestor.submobjects)
+            if mob in self.mobjects:
+                self.mobjects.remove(mob)
         return self
 
     def bring_to_front(self, *mobjects: Mobject):

manimlib/utils/family_ops.py

@@ -18,32 +18,3 @@ def extract_mobject_family_members(
         for sm in mob.get_family()
         if (not exclude_pointless) or sm.has_points()
     ]
-
-
-def restructure_list_to_exclude_certain_family_members(
-    mobject_list: list[Mobject],
-    to_remove: list[Mobject]
-) -> list[Mobject]:
-    """
-    Removes anything in to_remove from mobject_list, but in the event that one of
-    the items to be removed is a member of the family of an item in mobject_list,
-    the other family members are added back into the list.
-
-    This is useful in cases where a scene contains a group, e.g. Group(m1, m2, m3),
-    but one of its submobjects is removed, e.g. scene.remove(m1), it's useful
-    for the list of mobject_list to be edited to contain other submobjects, but not m1.
-    """
-    new_list = []
-    to_remove = extract_mobject_family_members(to_remove)
-
-    def add_safe_mobjects_from_list(list_to_examine, set_to_remove):
-        for mob in list_to_examine:
-            if mob in set_to_remove:
-                continue
-            intersect = set_to_remove.intersection(mob.get_family())
-            if intersect:
-                add_safe_mobjects_from_list(mob.submobjects, intersect)
-            else:
-                new_list.append(mob)
-    add_safe_mobjects_from_list(mobject_list, set(to_remove))
-    return new_list