Implemented VGroup parsing iterables (#3966)

manim/mobject/types/vectorized_mobject.py

@@ -24,6 +24,7 @@
 from manim.constants import *
 from manim.mobject.mobject import Mobject
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
+from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
 from manim.mobject.three_d.three_d_utils import (
     get_3d_vmob_gradient_start_and_end_points,
@@ -2056,7 +2057,11 @@ def construct(self):
 
     """
 
-    def __init__(self, *vmobjects, **kwargs):
+    def __init__(
+        self,
+        *vmobjects: VMobject | Iterable[VMobject],
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         self.add(*vmobjects)
 
@@ -2069,13 +2074,16 @@ def __str__(self) -> str:
             f"submobject{'s' if len(self.submobjects) > 0 else ''}"
         )
 
-    def add(self, *vmobjects: VMobject) -> Self:
-        """Checks if all passed elements are an instance of VMobject and then add them to submobjects
+    def add(
+        self,
+        *vmobjects: VMobject | Iterable[VMobject],
+    ) -> Self:
+        """Checks if all passed elements are an instance, or iterables of VMobject and then adds them to submobjects
 
         Parameters
         ----------
         vmobjects
-            List of VMobject to add
+            List or iterable of VMobjects to add
 
         Returns
         -------
@@ -2084,10 +2092,13 @@ def add(self, *vmobjects: VMobject) -> Self:
         Raises
         ------
         TypeError
-            If one element of the list is not an instance of VMobject
+            If one element of the list, or iterable is not an instance of VMobject
 
         Examples
         --------
+        The following example shows how to add individual or multiple `VMobject` instances through the `VGroup`
+        constructor and its `.add()` method.
+
         .. manim:: AddToVGroup
 
             class AddToVGroup(Scene):
@@ -2116,8 +2127,65 @@ def construct(self):
                     self.play( # Animate group without component
                         (gr-circle_red).animate.shift(RIGHT)
                     )
+
+        A `VGroup` can be created using iterables as well. Keep in mind that all generated values from an
+        iterable must be an instance of `VMobject`. This is demonstrated below:
+
+        .. manim:: AddIterableToVGroupExample
+            :save_last_frame:
+
+            class AddIterableToVGroupExample(Scene):
+                def construct(self):
+                    v = VGroup(
+                        Square(),               # Singular VMobject instance
+                        [Circle(), Triangle()], # List of VMobject instances
+                        Dot(),
+                        (Dot() for _ in range(2)), # Iterable that generates VMobjects
+                    )
+                    v.arrange()
+                    self.add(v)
+
+        To facilitate this, the iterable is unpacked before its individual instances are added to the `VGroup`.
+        As a result, when you index a `VGroup`, you will never get back an iterable.
+        Instead, you will always receive `VMobject` instances, including those
+        that were part of the iterable/s that you originally added to the `VGroup`.
         """
-        return super().add(*vmobjects)
+
+        def get_type_error_message(invalid_obj, invalid_indices):
+            return (
+                f"Only values of type {vmobject_render_type.__name__} can be added "
+                "as submobjects of VGroup, but the value "
+                f"{repr(invalid_obj)} (at index {invalid_indices[1]} of "
+                f"parameter {invalid_indices[0]}) is of type "
+                f"{type(invalid_obj).__name__}."
+            )
+
+        vmobject_render_type = (
+            OpenGLVMobject if config.renderer == RendererType.OPENGL else VMobject
+        )
+        valid_vmobjects = []
+
+        for i, vmobject in enumerate(vmobjects):
+            if isinstance(vmobject, vmobject_render_type):
+                valid_vmobjects.append(vmobject)
+            elif isinstance(vmobject, Iterable) and not isinstance(
+                vmobject, (Mobject, OpenGLMobject)
+            ):
+                for j, subvmobject in enumerate(vmobject):
+                    if not isinstance(subvmobject, vmobject_render_type):
+                        raise TypeError(get_type_error_message(subvmobject, (i, j)))
+                    valid_vmobjects.append(subvmobject)
+            elif isinstance(vmobject, Iterable) and isinstance(
+                vmobject, (Mobject, OpenGLMobject)
+            ):
+                raise TypeError(
+                    f"{get_type_error_message(vmobject, (i, 0))} "
+                    "You can try adding this value into a Group instead."
+                )
+            else:
+                raise TypeError(get_type_error_message(vmobject, (i, 0)))
+
+        return super().add(*valid_vmobjects)
 
     def __add__(self, vmobject: VMobject) -> Self:
         return VGroup(*self.submobjects, vmobject)

tests/module/mobject/types/vectorized_mobject/test_vectorized_mobject.py

@@ -132,26 +132,72 @@ def test_vgroup_init():
         VGroup(3.0)
     assert str(init_with_float_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value 3.0 (at index 0) is of type float."
+        "but the value 3.0 (at index 0 of parameter 0) is of type float."
     )
 
     with pytest.raises(TypeError) as init_with_mob_info:
         VGroup(Mobject())
     assert str(init_with_mob_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value Mobject (at index 0) is of type Mobject. You can try "
+        "but the value Mobject (at index 0 of parameter 0) is of type Mobject. You can try "
         "adding this value into a Group instead."
     )
 
     with pytest.raises(TypeError) as init_with_vmob_and_mob_info:
         VGroup(VMobject(), Mobject())
     assert str(init_with_vmob_and_mob_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value Mobject (at index 1) is of type Mobject. You can try "
+        "but the value Mobject (at index 0 of parameter 1) is of type Mobject. You can try "
         "adding this value into a Group instead."
     )
 
 
+def test_vgroup_init_with_iterable():
+    """Test VGroup instantiation with an iterable type."""
+
+    def type_generator(type_to_generate, n):
+        return (type_to_generate() for _ in range(n))
+
+    def mixed_type_generator(major_type, minor_type, minor_type_positions, n):
+        return (
+            minor_type() if i in minor_type_positions else major_type()
+            for i in range(n)
+        )
+
+    obj = VGroup(VMobject())
+    assert len(obj.submobjects) == 1
+
+    obj = VGroup(type_generator(VMobject, 38))
+    assert len(obj.submobjects) == 38
+
+    obj = VGroup(VMobject(), [VMobject(), VMobject()], type_generator(VMobject, 38))
+    assert len(obj.submobjects) == 41
+
+    # A VGroup cannot be initialised with an iterable containing a Mobject
+    with pytest.raises(TypeError) as init_with_mob_iterable:
+        VGroup(type_generator(Mobject, 5))
+    assert str(init_with_mob_iterable.value) == (
+        "Only values of type VMobject can be added as submobjects of VGroup, "
+        "but the value Mobject (at index 0 of parameter 0) is of type Mobject."
+    )
+
+    # A VGroup cannot be initialised with an iterable containing a Mobject in any position
+    with pytest.raises(TypeError) as init_with_mobs_and_vmobs_iterable:
+        VGroup(mixed_type_generator(VMobject, Mobject, [3, 5], 7))
+    assert str(init_with_mobs_and_vmobs_iterable.value) == (
+        "Only values of type VMobject can be added as submobjects of VGroup, "
+        "but the value Mobject (at index 3 of parameter 0) is of type Mobject."
+    )
+
+    # A VGroup cannot be initialised with an iterable containing non VMobject's in any position
+    with pytest.raises(TypeError) as init_with_float_and_vmobs_iterable:
+        VGroup(mixed_type_generator(VMobject, float, [6, 7], 9))
+    assert str(init_with_float_and_vmobs_iterable.value) == (
+        "Only values of type VMobject can be added as submobjects of VGroup, "
+        "but the value 0.0 (at index 6 of parameter 0) is of type float."
+    )
+
+
 def test_vgroup_add():
     """Test the VGroup add method."""
     obj = VGroup()
@@ -165,7 +211,7 @@ def test_vgroup_add():
         obj.add(3)
     assert str(add_int_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value 3 (at index 0) is of type int."
+        "but the value 3 (at index 0 of parameter 0) is of type int."
     )
     assert len(obj.submobjects) == 1
 
@@ -175,7 +221,7 @@ def test_vgroup_add():
         obj.add(Mobject())
     assert str(add_mob_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value Mobject (at index 0) is of type Mobject. You can try "
+        "but the value Mobject (at index 0 of parameter 0) is of type Mobject. You can try "
         "adding this value into a Group instead."
     )
     assert len(obj.submobjects) == 1
@@ -185,7 +231,7 @@ def test_vgroup_add():
         obj.add(VMobject(), Mobject())
     assert str(add_vmob_and_mob_info.value) == (
         "Only values of type VMobject can be added as submobjects of VGroup, "
-        "but the value Mobject (at index 1) is of type Mobject. You can try "
+        "but the value Mobject (at index 0 of parameter 1) is of type Mobject. You can try "
         "adding this value into a Group instead."
     )
     assert len(obj.submobjects) == 1

tests/opengl/test_opengl_vectorized_mobject.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from manim import Circle, Line, Square, VDict, VGroup
+from manim import Circle, Line, Square, VDict, VGroup, VMobject
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
 
@@ -90,26 +90,84 @@ def test_vgroup_init(using_opengl_renderer):
         VGroup(3.0)
     assert str(init_with_float_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value 3.0 (at index 0) is of type float."
+        "VGroup, but the value 3.0 (at index 0 of parameter 0) is of type float."
     )
 
     with pytest.raises(TypeError) as init_with_mob_info:
         VGroup(OpenGLMobject())
     assert str(init_with_mob_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value OpenGLMobject (at index 0) is of type "
+        "VGroup, but the value OpenGLMobject (at index 0 of parameter 0) is of type "
         "OpenGLMobject. You can try adding this value into a Group instead."
     )
 
     with pytest.raises(TypeError) as init_with_vmob_and_mob_info:
         VGroup(OpenGLVMobject(), OpenGLMobject())
     assert str(init_with_vmob_and_mob_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value OpenGLMobject (at index 1) is of type "
+        "VGroup, but the value OpenGLMobject (at index 0 of parameter 1) is of type "
         "OpenGLMobject. You can try adding this value into a Group instead."
     )
 
 
+def test_vgroup_init_with_iterable(using_opengl_renderer):
+    """Test VGroup instantiation with an iterable type."""
+
+    def type_generator(type_to_generate, n):
+        return (type_to_generate() for _ in range(n))
+
+    def mixed_type_generator(major_type, minor_type, minor_type_positions, n):
+        return (
+            minor_type() if i in minor_type_positions else major_type()
+            for i in range(n)
+        )
+
+    obj = VGroup(OpenGLVMobject())
+    assert len(obj.submobjects) == 1
+
+    obj = VGroup(type_generator(OpenGLVMobject, 38))
+    assert len(obj.submobjects) == 38
+
+    obj = VGroup(
+        OpenGLVMobject(),
+        [OpenGLVMobject(), OpenGLVMobject()],
+        type_generator(OpenGLVMobject, 38),
+    )
+    assert len(obj.submobjects) == 41
+
+    # A VGroup cannot be initialised with an iterable containing a OpenGLMobject
+    with pytest.raises(TypeError) as init_with_mob_iterable:
+        VGroup(type_generator(OpenGLMobject, 5))
+    assert str(init_with_mob_iterable.value) == (
+        "Only values of type OpenGLVMobject can be added as submobjects of VGroup, "
+        "but the value OpenGLMobject (at index 0 of parameter 0) is of type OpenGLMobject."
+    )
+
+    # A VGroup cannot be initialised with an iterable containing a OpenGLMobject in any position
+    with pytest.raises(TypeError) as init_with_mobs_and_vmobs_iterable:
+        VGroup(mixed_type_generator(OpenGLVMobject, OpenGLMobject, [3, 5], 7))
+    assert str(init_with_mobs_and_vmobs_iterable.value) == (
+        "Only values of type OpenGLVMobject can be added as submobjects of VGroup, "
+        "but the value OpenGLMobject (at index 3 of parameter 0) is of type OpenGLMobject."
+    )
+
+    # A VGroup cannot be initialised with an iterable containing non OpenGLVMobject's in any position
+    with pytest.raises(TypeError) as init_with_float_and_vmobs_iterable:
+        VGroup(mixed_type_generator(OpenGLVMobject, float, [6, 7], 9))
+    assert str(init_with_float_and_vmobs_iterable.value) == (
+        "Only values of type OpenGLVMobject can be added as submobjects of VGroup, "
+        "but the value 0.0 (at index 6 of parameter 0) is of type float."
+    )
+
+    # A VGroup cannot be initialised with an iterable containing both OpenGLVMobject's and VMobject's
+    with pytest.raises(TypeError) as init_with_mobs_and_vmobs_iterable:
+        VGroup(mixed_type_generator(OpenGLVMobject, VMobject, [3, 5], 7))
+    assert str(init_with_mobs_and_vmobs_iterable.value) == (
+        "Only values of type OpenGLVMobject can be added as submobjects of VGroup, "
+        "but the value VMobject (at index 3 of parameter 0) is of type VMobject."
+    )
+
+
 def test_vgroup_add(using_opengl_renderer):
     """Test the VGroup add method."""
     obj = VGroup()
@@ -123,7 +181,7 @@ def test_vgroup_add(using_opengl_renderer):
         obj.add(3)
     assert str(add_int_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value 3 (at index 0) is of type int."
+        "VGroup, but the value 3 (at index 0 of parameter 0) is of type int."
     )
     assert len(obj.submobjects) == 1
 
@@ -133,7 +191,7 @@ def test_vgroup_add(using_opengl_renderer):
         obj.add(OpenGLMobject())
     assert str(add_mob_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value OpenGLMobject (at index 0) is of type "
+        "VGroup, but the value OpenGLMobject (at index 0 of parameter 0) is of type "
         "OpenGLMobject. You can try adding this value into a Group instead."
     )
     assert len(obj.submobjects) == 1
@@ -143,7 +201,7 @@ def test_vgroup_add(using_opengl_renderer):
         obj.add(OpenGLVMobject(), OpenGLMobject())
     assert str(add_vmob_and_mob_info.value) == (
         "Only values of type OpenGLVMobject can be added as submobjects of "
-        "VGroup, but the value OpenGLMobject (at index 1) is of type "
+        "VGroup, but the value OpenGLMobject (at index 0 of parameter 1) is of type "
         "OpenGLMobject. You can try adding this value into a Group instead."
     )
     assert len(obj.submobjects) == 1