Add type annotations to manim.utils.* (#3999)

Co-authored-by: Francisco Manríquez Novoa <[email protected]>
Co-authored-by: Francisco Manríquez <[email protected]>

Author: henrikmidtiby

manim/mobject/text/numbers.py

@@ -5,6 +5,7 @@
 __all__ = ["DecimalNumber", "Integer", "Variable"]
 
 from collections.abc import Sequence
+from typing import Any
 
 import numpy as np
 
@@ -327,7 +328,9 @@ def construct(self):
                 self.add(Integer(number=6.28).set_x(-1.5).set_y(-2).set_color(YELLOW).scale(1.4))
     """
 
-    def __init__(self, number=0, num_decimal_places=0, **kwargs):
+    def __init__(
+        self, number: float = 0, num_decimal_places: int = 0, **kwargs: Any
+    ) -> None:
         super().__init__(number=number, num_decimal_places=num_decimal_places, **kwargs)
 
     def get_value(self):

manim/renderer/cairo_renderer.py

@@ -19,6 +19,8 @@
     from manim.animation.animation import Animation
     from manim.scene.scene import Scene
 
+    from ..typing import PixelArray
+
 __all__ = ["CairoRenderer"]
 
 
@@ -158,7 +160,7 @@ def render(self, scene, time, moving_mobjects):
         self.update_frame(scene, moving_mobjects)
         self.add_frame(self.get_frame())
 
-    def get_frame(self):
+    def get_frame(self) -> PixelArray:
         """
         Gets the current frame as NumPy array.
 

manim/renderer/opengl_renderer.py

@@ -388,7 +388,7 @@ def get_texture_id(self, path):
 
         return self.path_to_texture_id[repr(path)]
 
-    def update_skipping_status(self):
+    def update_skipping_status(self) -> None:
         """
         This method is used internally to check if the current
         animation needs to be skipped or not. It also checks if

manim/scene/scene.py

@@ -488,7 +488,7 @@ def add(self, *mobjects: Mobject):
                 self.moving_mobjects += mobjects
         return self
 
-    def add_mobjects_from_animations(self, animations):
+    def add_mobjects_from_animations(self, animations: list[Animation]) -> None:
         curr_mobjects = self.get_mobject_family_members()
         for animation in animations:
             if animation.is_introducer():

manim/scene/scene_file_writer.py

@@ -20,7 +20,7 @@
 from pydub import AudioSegment
 
 from manim import __version__
-from manim.typing import PixelArray
+from manim.typing import PixelArray, StrPath
 
 from .. import config, logger
 from .._config.logger_utils import set_file_logger
@@ -38,6 +38,7 @@
 from .section import DefaultSectionType, Section
 
 if TYPE_CHECKING:
+    from manim.renderer.cairo_renderer import CairoRenderer
     from manim.renderer.opengl_renderer import OpenGLRenderer
 
 
@@ -104,7 +105,12 @@ class SceneFileWriter:
 
     force_output_as_scene_name = False
 
-    def __init__(self, renderer, scene_name, **kwargs):
+    def __init__(
+        self,
+        renderer: CairoRenderer | OpenGLRenderer,
+        scene_name: StrPath,
+        **kwargs: Any,
+    ) -> None:
         self.renderer = renderer
         self.init_output_directories(scene_name)
         self.init_audio()
@@ -118,7 +124,7 @@ def __init__(self, renderer, scene_name, **kwargs):
             name="autocreated", type_=DefaultSectionType.NORMAL, skip_animations=False
         )
 
-    def init_output_directories(self, scene_name):
+    def init_output_directories(self, scene_name: StrPath) -> None:
         """Initialise output directories.
 
         Notes
@@ -378,7 +384,9 @@ def add_sound(
         self.add_audio_segment(new_segment, time, **kwargs)
 
     # Writers
-    def begin_animation(self, allow_write: bool = False, file_path=None):
+    def begin_animation(
+        self, allow_write: bool = False, file_path: StrPath | None = None
+    ) -> None:
         """
         Used internally by manim to stream the animation to FFMPEG for
         displaying or writing to a file.
@@ -391,7 +399,7 @@ def begin_animation(self, allow_write: bool = False, file_path=None):
         if write_to_movie() and allow_write:
             self.open_partial_movie_stream(file_path=file_path)
 
-    def end_animation(self, allow_write: bool = False):
+    def end_animation(self, allow_write: bool = False) -> None:
         """
         Internally used by Manim to stop streaming to
         FFMPEG gracefully.

manim/utils/bezier.py

@@ -29,15 +29,12 @@
 from manim.utils.simple_functions import choose
 
 if TYPE_CHECKING:
-    import numpy.typing as npt
-
     from manim.typing import (
         BezierPoints,
         BezierPoints_Array,
         BezierPointsLike,
         BezierPointsLike_Array,
         ColVector,
-        ManimFloat,
         MatrixMN,
         Point3D,
         Point3D_Array,
@@ -64,7 +61,9 @@ def bezier(
 ) -> Callable[[float | ColVector], Point3D_Array]: ...
 
 
-def bezier(points):
+def bezier(
+    points: Point3D_Array | Sequence[Point3D_Array],
+) -> Callable[[float | ColVector], Point3D_Array]:
     """Classic implementation of a Bézier curve.
 
     Parameters
@@ -118,21 +117,21 @@ def bezier(points):
 
     if degree == 0:
 
-        def zero_bezier(t):
+        def zero_bezier(t: float | ColVector) -> Point3D | Point3D_Array:
             return np.ones_like(t) * P[0]
 
         return zero_bezier
 
     if degree == 1:
 
-        def linear_bezier(t):
+        def linear_bezier(t: float | ColVector) -> Point3D | Point3D_Array:
             return P[0] + t * (P[1] - P[0])
 
         return linear_bezier
 
     if degree == 2:
 
-        def quadratic_bezier(t):
+        def quadratic_bezier(t: float | ColVector) -> Point3D | Point3D_Array:
             t2 = t * t
             mt = 1 - t
             mt2 = mt * mt
@@ -142,7 +141,7 @@ def quadratic_bezier(t):
 
     if degree == 3:
 
-        def cubic_bezier(t):
+        def cubic_bezier(t: float | ColVector) -> Point3D | Point3D_Array:
             t2 = t * t
             t3 = t2 * t
             mt = 1 - t
@@ -152,11 +151,12 @@ def cubic_bezier(t):
 
         return cubic_bezier
 
-    def nth_grade_bezier(t):
+    def nth_grade_bezier(t: float | ColVector) -> Point3D | Point3D_Array:
         is_scalar = not isinstance(t, np.ndarray)
         if is_scalar:
             B = np.empty((1, *P.shape))
         else:
+            assert isinstance(t, np.ndarray)
             t = t.reshape(-1, *[1 for dim in P.shape])
             B = np.empty((t.shape[0], *P.shape))
         B[:] = P
@@ -169,7 +169,8 @@ def nth_grade_bezier(t):
         # In the end, there shall be the evaluation at t of a single Bezier curve of
         # grade d, stored in the first slot of B
         if is_scalar:
-            return B[0, 0]
+            val: Point3D = B[0, 0]
+            return val
         return B[:, 0]
 
     return nth_grade_bezier
@@ -1026,7 +1027,11 @@ def interpolate(start: Point3D, end: Point3D, alpha: float) -> Point3D: ...
 def interpolate(start: Point3D, end: Point3D, alpha: ColVector) -> Point3D_Array: ...
 
 
-def interpolate(start, end, alpha):
+def interpolate(
+    start: float | Point3D,
+    end: float | Point3D,
+    alpha: float | ColVector,
+) -> float | ColVector | Point3D | Point3D_Array:
     """Linearly interpolates between two values ``start`` and ``end``.
 
     Parameters
@@ -1139,7 +1144,9 @@ def inverse_interpolate(start: Point3D, end: Point3D, value: Point3D) -> Point3D
 
 
 def inverse_interpolate(
-    start: float | Point3D, end: float | Point3D, value: float | Point3D
+    start: float | Point3D,
+    end: float | Point3D,
+    value: float | Point3D,
 ) -> float | Point3D:
     """Perform inverse interpolation to determine the alpha
     values that would produce the specified ``value``
@@ -1234,7 +1241,7 @@ def match_interpolate(
     return interpolate(
         new_start,
         new_end,
-        old_alpha,  # type: ignore[arg-type]
+        old_alpha,
     )
 
 
@@ -1270,7 +1277,8 @@ def get_smooth_cubic_bezier_handle_points(
     # they can only be an interpolation of these two anchors with alphas
     # 1/3 and 2/3, which will draw a straight line between the anchors.
     if n_anchors == 2:
-        return interpolate(anchors[0], anchors[1], np.array([[1 / 3], [2 / 3]]))
+        val = interpolate(anchors[0], anchors[1], np.array([[1 / 3], [2 / 3]]))
+        return (val[0], val[1])
 
     # Handle different cases depending on whether the points form a closed
     # curve or not
@@ -1745,7 +1753,12 @@ def get_quadratic_approximation_of_cubic(
 ) -> QuadraticBezierPath: ...
 
 
-def get_quadratic_approximation_of_cubic(a0, h0, h1, a1):
+def get_quadratic_approximation_of_cubic(
+    a0: Point3D | Point3D_Array,
+    h0: Point3D | Point3D_Array,
+    h1: Point3D | Point3D_Array,
+    a1: Point3D | Point3D_Array,
+) -> QuadraticSpline | QuadraticBezierPath:
     r"""If ``a0``, ``h0``, ``h1`` and ``a1`` are the control points of a cubic
     Bézier curve, approximate the curve with two quadratic Bézier curves and
     return an array of 6 points, where the first 3 points represent the first
@@ -1849,33 +1862,33 @@ def get_quadratic_approximation_of_cubic(a0, h0, h1, a1):
         If ``a0``, ``h0``, ``h1`` and ``a1`` have different dimensions, or
         if their number of dimensions is not 1 or 2.
     """
-    a0 = np.asarray(a0)
-    h0 = np.asarray(h0)
-    h1 = np.asarray(h1)
-    a1 = np.asarray(a1)
-
-    if all(arr.ndim == 1 for arr in (a0, h0, h1, a1)):
-        num_curves, dim = 1, a0.shape[0]
-    elif all(arr.ndim == 2 for arr in (a0, h0, h1, a1)):
-        num_curves, dim = a0.shape
+    a0c = np.asarray(a0)
+    h0c = np.asarray(h0)
+    h1c = np.asarray(h1)
+    a1c = np.asarray(a1)
+
+    if all(arr.ndim == 1 for arr in (a0c, h0c, h1c, a1c)):
+        num_curves, dim = 1, a0c.shape[0]
+    elif all(arr.ndim == 2 for arr in (a0c, h0c, h1c, a1c)):
+        num_curves, dim = a0c.shape
     else:
         raise ValueError("All arguments must be Point3D or Point3D_Array.")
 
-    m0 = 0.25 * (3 * h0 + a0)
-    m1 = 0.25 * (3 * h1 + a1)
+    m0 = 0.25 * (3 * h0c + a0c)
+    m1 = 0.25 * (3 * h1c + a1c)
     k = 0.5 * (m0 + m1)
 
     result = np.empty((6 * num_curves, dim))
-    result[0::6] = a0
+    result[0::6] = a0c
     result[1::6] = m0
     result[2::6] = k
     result[3::6] = k
     result[4::6] = m1
-    result[5::6] = a1
+    result[5::6] = a1c
     return result
 
 
-def is_closed(points: Point3DLike_Array) -> bool:
+def is_closed(points: Point3D_Array) -> bool:
     """Returns ``True`` if the spline given by ``points`` is closed, by
     checking if its first and last points are close to each other, or``False``
     otherwise.
@@ -1952,7 +1965,7 @@ def proportions_along_bezier_curve_for_point(
     point: Point3DLike,
     control_points: BezierPointsLike,
     round_to: float = 1e-6,
-) -> npt.NDArray[ManimFloat]:
+) -> MatrixMN:
     """Obtains the proportion along the bezier curve corresponding to a given point
     given the bezier curve's control points.
 

manim/utils/caching.py

@@ -1,14 +1,20 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 from .. import config, logger
 from ..utils.hashing import get_hash_from_play_call
 
 __all__ = ["handle_caching_play"]
 
+if TYPE_CHECKING:
+    from typing import Any
 
-def handle_caching_play(func: Callable[..., None]):
+    from manim.renderer.opengl_renderer import OpenGLRenderer
+    from manim.scene.scene import Scene
+
+
+def handle_caching_play(func: Callable[..., None]) -> Callable[..., None]:
     """Decorator that returns a wrapped version of func that will compute
     the hash of the play invocation.
 
@@ -28,7 +34,7 @@ def handle_caching_play(func: Callable[..., None]):
     # the play logic of the latter has to be refactored in the same way the cairo renderer has been, and thus this
     # method has to be deleted.
 
-    def wrapper(self, scene, *args, **kwargs):
+    def wrapper(self: OpenGLRenderer, scene: Scene, *args: Any, **kwargs: Any) -> None:
         self.skip_animations = self._original_skipping_status
         self.update_skipping_status()
         animations = scene.compile_animations(*args, **kwargs)
@@ -43,8 +49,9 @@ def wrapper(self, scene, *args, **kwargs):
             return
         if not config["disable_caching"]:
             mobjects_on_scene = scene.mobjects
+            # TODO: the first argument seems wrong. Shouldn't it be scene instead?
             hash_play = get_hash_from_play_call(
-                self,
+                self,  # type: ignore[arg-type]
                 self.camera,
                 animations,
                 mobjects_on_scene,