Added docs for functions in mobject_update_utils

manim/animation/updaters/mobject_update_utils.py

@@ -15,30 +15,34 @@
 
 
 import inspect
-from collections.abc import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
 
 from manim.constants import DEGREES, RIGHT
 from manim.mobject.mobject import Mobject
 from manim.opengl import OpenGLMobject
+from manim.utils.space_ops import normalize
 
+if TYPE_CHECKING:
+    from manim.animation.animation import Animation
 
-def assert_is_mobject_method(method):
+
+def assert_is_mobject_method(method: Callable) -> None:
     assert inspect.ismethod(method)
     mobject = method.__self__
     assert isinstance(mobject, (Mobject, OpenGLMobject))
 
 
-def always(method, *args, **kwargs):
+def always(method: Callable, *args, **kwargs) -> Mobject:
     assert_is_mobject_method(method)
     mobject = method.__self__
     func = method.__func__
     mobject.add_updater(lambda m: func(m, *args, **kwargs))
     return mobject
 
 
-def f_always(method, *arg_generators, **kwargs):
+def f_always(method: Callable[[Mobject], None], *arg_generators, **kwargs) -> Mobject:
     """
     More functional version of always, where instead
     of taking in args, it takes in functions which output
@@ -80,16 +84,18 @@ def construct(self):
                 sine = ax.plot(np.sin, color=RED)
                 alpha = ValueTracker(0)
                 point = always_redraw(
-                        lambda: Dot(
-                            sine.point_from_proportion(alpha.get_value()),
-                            color=BLUE)
+                    lambda: Dot(
+                        sine.point_from_proportion(alpha.get_value()),
+                        color=BLUE
                     )
+                )
                 tangent = always_redraw(
                     lambda: TangentLine(
                         sine,
                         alpha=alpha.get_value(),
                         color=YELLOW,
-                        length=4)
+                        length=4
+                    )
                 )
                 self.add(ax, sine, point, tangent)
                 self.play(alpha.animate.set_value(1), rate_func=linear, run_time=2)
@@ -99,36 +105,110 @@ def construct(self):
     return mob
 
 
-def always_shift(mobject, direction=RIGHT, rate=0.1):
-    def normalize(v):
-        norm = np.linalg.norm(v)
-        if norm == 0:
-            return v
-        return v / norm
+def always_shift(
+    mobject: Mobject, direction: np.ndarray[np.float64] = RIGHT, rate: float = 0.1
+) -> Mobject:
+    """A mobject which is continuously shifted along some direction
+    at a certain rate.
+
+    Parameters
+    ----------
+    mobject
+        The mobject to shift.
+    direction
+        The direction to shift. The vector is normalized, the specified magnitude
+        is not relevant.
+    rate
+        Length in Manim units which the mobject travels in one
+        second along the specified direction.
+
+    Examples
+    --------
+
+    .. manim:: ShiftingSquare
+
+        class ShiftingSquare(Scene):
+            def construct(self):
+                sq = Square().set_fill(opacity=1)
+                tri = Triangle()
+                VGroup(sq, tri).arrange(LEFT)
+
+                # construct a square which is continuously
+                # shifted to the right
+                always_shift(sq, RIGHT, rate=5)
 
+                self.add(sq)
+                self.play(tri.animate.set_fill(opacity=1))
+    """
     mobject.add_updater(lambda m, dt: m.shift(dt * rate * normalize(direction)))
     return mobject
 
 
-def always_rotate(mobject, rate=20 * DEGREES, **kwargs):
+def always_rotate(mobject: Mobject, rate: float = 20 * DEGREES, **kwargs) -> Mobject:
+    """A mobject which is continuously rotated at a certain rate.
+
+    Parameters
+    ----------
+    mobject
+        The mobject to be rotated.
+    rate
+        The angle which the mobject is rotated by
+        over one second.
+    kwags
+        Further arguments to be passed to :meth:`.Mobject.rotate`.
+
+    Examples
+    --------
+
+    .. manim:: SpinningTriangle
+
+        class SpinningTriangle(Scene):
+            def construct(self):
+                tri = Triangle().set_fill(opacity=1).set_z_index(2)
+                sq = Square().to_edge(LEFT)
+
+                # will keep spinning while there is an animation going on
+                always_rotate(tri, rate=2*PI, about_point=ORIGIN)
+
+                self.add(tri, sq)
+                self.play(sq.animate.to_edge(RIGHT), rate_func=linear, run_time=1)
+    """
     mobject.add_updater(lambda m, dt: m.rotate(dt * rate, **kwargs))
     return mobject
 
 
-def turn_animation_into_updater(animation, cycle=False, **kwargs):
+def turn_animation_into_updater(
+    animation: Animation, cycle: bool = False, **kwargs
+) -> Mobject:
     """
     Add an updater to the animation's mobject which applies
     the interpolation and update functions of the animation
 
     If cycle is True, this repeats over and over.  Otherwise,
     the updater will be popped upon completion
+
+    Examples
+    --------
+
+    .. manim:: WelcomeToManim
+
+        class WelcomeToManim(Scene):
+            def construct(self):
+                words = Text("Welcome to")
+                banner = ManimBanner().scale(0.5)
+                VGroup(words, banner).arrange(DOWN)
+
+                turn_animation_into_updater(Write(words, run_time=0.9))
+                self.add(words)
+                self.wait(0.5)
+                self.play(banner.expand(), run_time=0.5)
     """
     mobject = animation.mobject
     animation.suspend_mobject_updating = False
     animation.begin()
     animation.total_time = 0
 
-    def update(m, dt):
+    def update(m: Mobject, dt: float):
         run_time = animation.get_run_time()
         time_ratio = animation.total_time / run_time
         if cycle:
@@ -147,5 +227,5 @@ def update(m, dt):
     return mobject
 
 
-def cycle_animation(animation, **kwargs):
+def cycle_animation(animation: Animation, **kwargs) -> Mobject:
     return turn_animation_into_updater(animation, cycle=True, **kwargs)