Separate Animation and Scene using SceneBuffer

Author: JasonGrace2282

manim/animation/animation.py

@@ -11,15 +11,17 @@
 from ..mobject.mobject import Mobject
 from ..mobject.opengl import opengl_mobject
 from ..utils.rate_functions import linear, smooth
+from .scene_buffer import SceneBuffer
+from .protocol import AnimationProtocol
 
 __all__ = ["Animation", "Wait", "override_animation"]
 
 
 from copy import deepcopy
-from typing import TYPE_CHECKING, Callable, Iterable, Sequence
+from typing import TYPE_CHECKING, Callable, Iterable, TypeVar, Sequence
 
 if TYPE_CHECKING:
-    from manim.scene.scene import Scene
+    from typing_extensions import Self
 
 
 DEFAULT_ANIMATION_RUN_TIME: float = 1.0
@@ -127,17 +129,17 @@ def __new__(
 
     def __init__(
         self,
-        mobject: Mobject | None,
+        mobject: OpenGLMobject | None,
         lag_ratio: float = DEFAULT_ANIMATION_LAG_RATIO,
         run_time: float = DEFAULT_ANIMATION_RUN_TIME,
         rate_func: Callable[[float], float] = smooth,
         reverse_rate_function: bool = False,
-        name: str = None,
-        remover: bool = False,  # remove a mobject from the screen?
+        name: str | None = None,
+        remover: bool = False,  # remove a mobject from the screen at end of animation
         suspend_mobject_updating: bool = True,
         introducer: bool = False,
         *,
-        _on_finish: Callable[[], None] = lambda _: None,
+        _on_finish: Callable[[SceneBuffer], object] = lambda _: None,
         **kwargs,
     ) -> None:
         self._typecheck_input(mobject)
@@ -149,15 +151,16 @@ def __init__(
         self.introducer: bool = introducer
         self.suspend_mobject_updating: bool = suspend_mobject_updating
         self.lag_ratio: float = lag_ratio
-        self._on_finish: Callable[[Scene], None] = _on_finish
+        self._on_finish = _on_finish
+        self.buffer = SceneBuffer()
         if config["renderer"] == RendererType.OPENGL:
             self.starting_mobject: OpenGLMobject = OpenGLMobject()
             self.mobject: OpenGLMobject = (
                 mobject if mobject is not None else OpenGLMobject()
             )
-        else:
-            self.starting_mobject: Mobject = Mobject()
-            self.mobject: Mobject = mobject if mobject is not None else Mobject()
+        # else:
+        #     self.starting_mobject: Mobject = Mobject()
+        #     self.mobject: Mobject = mobject if mobject is not None else Mobject()
         if kwargs:
             logger.debug("Animation received extra kwargs: %s", kwargs)
 
@@ -169,7 +172,7 @@ def __init__(
                 ),
             )
 
-    def _typecheck_input(self, mobject: Mobject | None) -> None:
+    def _typecheck_input(self, mobject: Mobject | OpenGLMobject | None) -> None:
         if mobject is None:
             logger.debug("Animation with empty mobject")
         elif not isinstance(mobject, (Mobject, OpenGLMobject)):
@@ -213,10 +216,10 @@ def begin(self) -> None:
             self.mobject.suspend_updating()
         self.interpolate(0)
 
+        if self.is_introducer():
+            self.buffer.add(self.mobject)
+
     def finish(self) -> None:
-        # TODO: begin and finish should require a scene as parameter.
-        # That way Animation.clean_up_from_screen and Scene.add_mobjects_from_animations
-        # could be removed as they fulfill basically the same purpose.
         """Finish the animation.
 
         This method gets called when the animation is over.
@@ -226,44 +229,20 @@ def finish(self) -> None:
         if self.suspend_mobject_updating and self.mobject is not None:
             self.mobject.resume_updating()
 
-    def clean_up_from_scene(self, scene: Scene) -> None:
-        """Clean up the :class:`~.Scene` after finishing the animation.
-
-        This includes to :meth:`~.Scene.remove` the Animation's
-        :class:`~.Mobject` if the animation is a remover.
-
-        Parameters
-        ----------
-        scene
-            The scene the animation should be cleaned up from.
-        """
-        self._on_finish(scene)
+        self._on_finish(self.buffer)
         if self.is_remover():
-            scene.remove(self.mobject)
-
-    def _setup_scene(self, scene: Scene) -> None:
-        """Setup up the :class:`~.Scene` before starting the animation.
-
-        This includes to :meth:`~.Scene.add` the Animation's
-        :class:`~.Mobject` if the animation is an introducer.
-
-        Parameters
-        ----------
-        scene
-            The scene the animation should be cleaned up from.
-        """
-        if scene is None:
-            return
-        if (
-            self.is_introducer()
-            and self.mobject not in scene.get_mobject_family_members()
-        ):
-            scene.add(self.mobject)
+            self.buffer.remove(self.mobject)
 
     def create_starting_mobject(self) -> Mobject:
         # Keep track of where the mobject starts
         return self.mobject.copy()
 
+    def get_all_animations(self) -> tuple[Animation, ...]:
+        """This method is to implement an animation protocol, and
+        is more useful in places like :class:`.AnimationGroup`
+        """
+        return (self,)
+
     def get_all_mobjects(self) -> Sequence[Mobject]:
         """Get all mobjects involved in the animation.
 
@@ -305,9 +284,9 @@ def get_all_mobjects_to_update(self) -> list[Mobject]:
         # The surrounding scene typically handles
         # updating of self.mobject.  Besides, in
         # most cases its updating is suspended anyway
-        return list(filter(lambda m: m is not self.mobject, self.get_all_mobjects()))
+        return [m for m in self.get_all_mobjects() if m is not self.mobject]
 
-    def copy(self) -> Animation:
+    def copy(self) -> Self:
         """Create a copy of the animation.
 
         Returns
@@ -422,7 +401,7 @@ def get_run_time(self) -> float:
     def set_rate_func(
         self,
         rate_func: Callable[[float], float],
-    ) -> Animation:
+    ) -> Self:
         """Set the rate function of the animation.
 
         Parameters
@@ -451,7 +430,7 @@ def get_rate_func(
         """
         return self.rate_func
 
-    def set_name(self, name: str) -> Animation:
+    def set_name(self, name: str) -> Self:
         """Set the name of the animation.
 
         Parameters
@@ -489,7 +468,7 @@ def is_introducer(self) -> bool:
 
 
 def prepare_animation(
-    anim: Animation | mobject._AnimationBuilder,
+    anim: AnimationProtocol | mobject._AnimationBuilder | opengl_mobject._AnimationBuilder,
 ) -> Animation:
     r"""Returns either an unchanged animation, or the animation built
     from a passed animation factory.
@@ -517,10 +496,7 @@ def prepare_animation(
         TypeError: Object 42 cannot be converted to an animation
 
     """
-    if isinstance(anim, mobject._AnimationBuilder):
-        return anim.build()
-
-    if isinstance(anim, opengl_mobject._AnimationBuilder):
+    if isinstance(anim, (mobject._AnimationBuilder, opengl_mobject._AnimationBuilder)):
         return anim.build()
 
     if isinstance(anim, Animation):
@@ -576,7 +552,7 @@ def begin(self) -> None:
     def finish(self) -> None:
         pass
 
-    def clean_up_from_scene(self, scene: Scene) -> None:
+    def clean_up_from_scene(self, scene: SceneBuffer) -> None:
         pass
 
     def update_mobjects(self, dt: float) -> None:
@@ -625,7 +601,9 @@ def construct(self):
 
     """
 
-    def decorator(func):
+    _F = TypeVar("_F", bound=Callable)
+
+    def decorator(func: _F) -> _F:
         func._override_animation = animation_class
         return func
 

manim/animation/protocol.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, Sequence
+
+if TYPE_CHECKING:
+    from .animation import Animation
+    from .scene_buffer import SceneBuffer
+
+
+__all__ = (
+    "AnimationProtocol",
+)
+
+
+class AnimationProtocol(Protocol):
+    buffer: SceneBuffer
+
+    def begin(self) -> None:
+        ...
+
+    def finish(self) -> None:
+        ...
+
+    def get_all_animations(self) -> Sequence[Animation]:
+        ...
+
+    def update_mobjects(self, dt: float) -> None:
+        ...
+
+    def interpolate(self, alpha: float) -> None:
+        ...
+
+    def get_run_time(self) -> float:
+        ...

manim/animation/scene_buffer.py

@@ -0,0 +1,65 @@
+from typing import TYPE_CHECKING, final, Sequence
+
+if TYPE_CHECKING:
+    from manim.mobject.opengl.opengl_mobject import OpenGLMobject as Mobject
+
+
+__all__ = ["SceneBuffer"]
+
+
+@final
+class SceneBuffer:
+    """
+    A "buffer" between :class:`.Scene` and :class:`.Animation`
+
+    Operations an animation wants to do on :class:`.Scene` should be
+    done here (eg. :meth:`.Scene.add`, :meth:`.Scene.remove`). The
+    scene will then apply these changes at specific points (namely
+    at the beginning and end of animations)
+
+    It is the scenes job to clear the buffer in between the beginning
+    and end of animations.
+    """
+
+    def __init__(self) -> None:
+        self.to_remove: list[Mobject] = []
+        self.to_add: list[Mobject] = []
+        self.deferred = False
+
+    def add(self, *mobs: Mobject) -> None:
+        check_args(mobs)
+        self._check_deferred()
+        self.to_add.extend(mobs)
+
+    def remove(self, *mobs: Mobject) -> None:
+        check_args(mobs)
+        self._check_deferred()
+        self.to_remove.extend(mobs)
+
+    def clear(self) -> None:
+        self.to_remove.clear()
+        self.to_add.clear()
+
+    def deferred_clear(self) -> None:
+        """Clear ``self`` on next operation"""
+        self.deferred = True
+
+    def _check_deferred(self) -> None:
+        if self.deferred:
+            self.clear()
+            self.deferred = False
+
+    def __str__(self) -> str:
+        to_add = self.to_add
+        to_remove = self.to_remove
+        return f"{type(self).__name__}({to_add=}, {to_remove=})"
+
+    __repr__ = __str__
+
+
+def check_args(mobs: Sequence[Mobject]) -> None:
+    if not (
+        isinstance(mobs, Sequence)
+        and all(isinstance(x, Mobject) for x in mobs)
+    ):
+        raise TypeError("Must set to a list of Mobjects")