feat: implement __add__, __iadd__, __sub__, and __isub__ for Mobject

Co-authored-by: Leo Torres <[email protected]>

Author: jsonvillanueva

docs/source/changelog.rst

@@ -93,6 +93,7 @@ Mobjects, Scenes, and Animations
 #. Added BraceBetweenPoints (via :pr:`693`).
 #. Added ArcPolygon and ArcPolygonFromArcs (via :pr:`707`).
 #. Added Cutout (via :pr:`760`).
+#. Added Mobject raise not implemented errors for dunder methods and implementations for VGroup dunder methods (via :pr:`790`).
 #. Added :class:`~.ManimBanner` for a animated version of our logo and banner (via :pr:`729`)
 #. The background color of a scene can now be changed reliably by setting, e.g.,
    ``self.camera.background_color = RED`` (via :pr:`716`).

manim/mobject/mobject.py

@@ -77,6 +77,8 @@ def add(self, *mobjects):
 
         The mobjects are added to self.submobjects.
 
+        Subclasses of mobject may implement + and += dunder methods.
+
         Parameters
         ----------
         mobjects : :class:`Mobject`
@@ -136,6 +138,12 @@ def add(self, *mobjects):
         self.submobjects = list_update(self.submobjects, mobjects)
         return self
 
+    def __add__(self, mobject):
+        raise NotImplementedError
+
+    def __iadd__(self, mobject):
+        raise NotImplementedError
+
     def add_to_back(self, *mobjects):
         self.remove(*mobjects)
         self.submobjects = list(mobjects) + self.submobjects
@@ -146,6 +154,8 @@ def remove(self, *mobjects):
 
         The mobjects are removed from self.submobjects, if they exist.
 
+        Subclasses of mobject may implement - and -= dunder methods.
+
         Parameters
         ----------
         mobjects : :class:`Mobject`
@@ -166,6 +176,12 @@ def remove(self, *mobjects):
                 self.submobjects.remove(mobject)
         return self
 
+    def __sub__(self, other):
+        raise NotImplementedError
+
+    def __isub__(self, other):
+        raise NotImplementedError
+
     def get_array_attrs(self):
         return ["points"]
 

manim/mobject/types/vectorized_mobject.py

@@ -975,6 +975,31 @@ class VGroup(VMobject):
     Examples
     --------
 
+    To add :class:`~.VMobject`s to a :class:`~.VGroup`, you can either use the
+    :meth:`~.VGroup.add` method, or use the `+` and `+=` operators. Similarly, you
+    can subtract elements of a VGroup via :meth:`~.VGroup.remove` method, or
+    `-` and `-=` operators:
+
+        >>> from manim import Triangle, Square, VGroup
+        >>> vg = VGroup()
+        >>> triangle, square = Triangle(), Square()
+        >>> vg.add(triangle)
+        VGroup(Triangle)
+        >>> vg + square   # a new VGroup is constructed
+        VGroup(Triangle, Square)
+        >>> vg            # not modified
+        VGroup(Triangle)
+        >>> vg += square; vg  # modifies vg
+        VGroup(Triangle, Square)
+        >>> vg.remove(triangle)
+        VGroup(Square)
+        >>> vg - square; # a new VGroup is constructed
+        VGroup()
+        >>> vg   # not modified
+        VGroup(Square)
+        >>> vg -= square; vg # modifies vg
+        VGroup()
+
     .. manim:: ArcShapeIris
         :save_last_frame:
 
@@ -1014,16 +1039,61 @@ def add(self, *vmobjects):
 
         Returns
         -------
-        None
+        :class:`VGroup`
 
         Raises
         ------
         TypeError
             If one element of the list is not an instance of VMobject
+
+        Examples
+        --------
+        .. manim:: AddToVGroup
+
+            class AddToVGroup(Scene):
+                def construct(self):
+                    circle_red = Circle(color=RED)
+                    circle_green = Circle(color=GREEN)
+                    circle_blue = Circle(color=BLUE)
+                    circle_red.shift(LEFT)
+                    circle_blue.shift(RIGHT)
+                    gr = VGroup(circle_red, circle_green)
+                    gr2 = VGroup(circle_blue) # Constructor uses add directly
+                    self.add(gr,gr2)
+                    self.wait()
+                    gr += gr2 # Add group to another
+                    self.play(
+                        gr.shift, DOWN,
+                    )
+                    gr -= gr2 # Remove group
+                    self.play( # Animate groups separately
+                        gr.shift, LEFT,
+                        gr2.shift, UP,
+                    )
+                    self.play( #Animate groups without modification
+                        (gr+gr2).shift, RIGHT
+                    )
+                    self.play( # Animate group without component
+                        (gr-circle_red).shift, RIGHT
+                    )
         """
         if not all(isinstance(m, VMobject) for m in vmobjects):
             raise TypeError("All submobjects must be of type VMobject")
-        super().add(*vmobjects)
+        return super().add(*vmobjects)
+
+    def __add__(self, vmobject):
+        return VGroup(*self.submobjects, vmobject)
+
+    def __iadd__(self, vmobject):
+        return self.add(vmobject)
+
+    def __sub__(self, vmobject):
+        copy = VGroup(*self.submobjects)
+        copy.remove(vmobject)
+        return copy
+
+    def __isub__(self, vmobject):
+        return self.remove(vmobject)
 
 
 class VDict(VMobject):

tests/test_vectorized_mobject.py

@@ -31,6 +31,60 @@ def test_vgroup_add():
         obj.add(obj)
 
 
+def test_vgroup_add_dunder():
+    """Test the VGroup __add__ magic method."""
+    obj = VGroup()
+    assert len(obj.submobjects) == 0
+    obj + VMobject()
+    assert len(obj.submobjects) == 0
+    obj += VMobject()
+    assert len(obj.submobjects) == 1
+    with pytest.raises(TypeError):
+        obj += Mobject()
+    assert len(obj.submobjects) == 1
+    with pytest.raises(TypeError):
+        # If only one of the added object is not an instance of VMobject, none of them should be added
+        obj += (VMobject(), Mobject())
+    assert len(obj.submobjects) == 1
+    with pytest.raises(Exception):  # TODO change this to ValueError once #307 is merged
+        # a Mobject cannot contain itself
+        obj += obj
+
+
+def test_vgroup_remove():
+    """Test the VGroup remove method."""
+    a = VMobject()
+    c = VMobject()
+    b = VGroup(c)
+    obj = VGroup(a, b)
+    assert len(obj.submobjects) == 2
+    assert len(b.submobjects) == 1
+    obj.remove(a)
+    b.remove(c)
+    assert len(obj.submobjects) == 1
+    assert len(b.submobjects) == 0
+    obj.remove(b)
+    assert len(obj.submobjects) == 0
+
+
+def test_vgroup_remove_dunder():
+    """Test the VGroup __sub__ magic method."""
+    a = VMobject()
+    c = VMobject()
+    b = VGroup(c)
+    obj = VGroup(a, b)
+    assert len(obj.submobjects) == 2
+    assert len(b.submobjects) == 1
+    assert len((obj - a)) == 1
+    assert len(obj.submobjects) == 2
+    obj -= a
+    b -= c
+    assert len(obj.submobjects) == 1
+    assert len(b.submobjects) == 0
+    obj -= b
+    assert len(obj.submobjects) == 0
+
+
 def test_vdict_init():
     """Test the VDict instantiation."""
     # Test empty VDict