Add type annotations to dot_cloud.py, vectorized_mobject_rendering.py and opengl_three_dimensions.py (#4359)

Co-authored-by: Francisco Manríquez Novoa <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

manim/mobject/opengl/dot_cloud.py

@@ -2,24 +2,33 @@
 
 __all__ = ["TrueDot", "DotCloud"]
 
+from typing import Any
+
 import numpy as np
+from typing_extensions import Self
 
 from manim.constants import ORIGIN, RIGHT, UP
 from manim.mobject.opengl.opengl_point_cloud_mobject import OpenGLPMobject
-from manim.utils.color import YELLOW
+from manim.typing import Point3DLike
+from manim.utils.color import YELLOW, ParsableManimColor
 
 
 class DotCloud(OpenGLPMobject):
     def __init__(
-        self, color=YELLOW, stroke_width=2.0, radius=2.0, density=10, **kwargs
+        self,
+        color: ParsableManimColor = YELLOW,
+        stroke_width: float = 2.0,
+        radius: float = 2.0,
+        density: float = 10,
+        **kwargs: Any,
     ):
         self.radius = radius
         self.epsilon = 1.0 / density
         super().__init__(
             stroke_width=stroke_width, density=density, color=color, **kwargs
         )
 
-    def init_points(self):
+    def init_points(self) -> None:
         self.points = np.array(
             [
                 r * (np.cos(theta) * RIGHT + np.sin(theta) * UP)
@@ -34,14 +43,16 @@ def init_points(self):
             dtype=np.float32,
         )
 
-    def make_3d(self, gloss=0.5, shadow=0.2):
+    def make_3d(self, gloss: float = 0.5, shadow: float = 0.2) -> Self:
         self.set_gloss(gloss)
         self.set_shadow(shadow)
         self.apply_depth_test()
         return self
 
 
 class TrueDot(DotCloud):
-    def __init__(self, center=ORIGIN, stroke_width=2.0, **kwargs):
+    def __init__(
+        self, center: Point3DLike = ORIGIN, stroke_width: float = 2.0, **kwargs: Any
+    ):
         self.radius = stroke_width
         super().__init__(points=[center], stroke_width=stroke_width, **kwargs)

manim/mobject/opengl/opengl_point_cloud_mobject.py

@@ -4,11 +4,19 @@
 
 import moderngl
 import numpy as np
+from typing_extensions import Self
 
 from manim.constants import *
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.utils.bezier import interpolate
-from manim.utils.color import BLACK, WHITE, YELLOW, color_gradient, color_to_rgba
+from manim.utils.color import (
+    BLACK,
+    WHITE,
+    YELLOW,
+    ParsableManimColor,
+    color_gradient,
+    color_to_rgba,
+)
 from manim.utils.config_ops import _Uniforms
 from manim.utils.iterables import resize_with_interpolation
 
@@ -27,15 +35,19 @@ class OpenGLPMobject(OpenGLMobject):
     point_radius = _Uniforms()
 
     def __init__(
-        self, stroke_width=2.0, color=YELLOW, render_primitive=moderngl.POINTS, **kwargs
+        self,
+        stroke_width: float = 2.0,
+        color: ParsableManimColor = YELLOW,
+        render_primitive: int = moderngl.POINTS,
+        **kwargs,
     ):
         self.stroke_width = stroke_width
         super().__init__(color=color, render_primitive=render_primitive, **kwargs)
         self.point_radius = (
             self.stroke_width * OpenGLPMobject.OPENGL_POINT_RADIUS_SCALE_FACTOR
         )
 
-    def reset_points(self):
+    def reset_points(self) -> Self:
         self.rgbas = np.zeros((1, 4))
         self.points = np.zeros((0, 3))
         return self

manim/mobject/opengl/opengl_surface.py

@@ -9,6 +9,7 @@
 
 from manim.constants import *
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
+from manim.typing import Point3D_Array, Vector3D_Array
 from manim.utils.bezier import integer_interpolate, interpolate
 from manim.utils.color import *
 from manim.utils.config_ops import _Data, _Uniforms
@@ -160,12 +161,14 @@ def compute_triangle_indices(self):
     def get_triangle_indices(self):
         return self.triangle_indices
 
-    def get_surface_points_and_nudged_points(self):
+    def get_surface_points_and_nudged_points(
+        self,
+    ) -> tuple[Point3D_Array, Point3D_Array, Point3D_Array]:
         points = self.points
         k = len(points) // 3
         return points[:k], points[k : 2 * k], points[2 * k :]
 
-    def get_unit_normals(self):
+    def get_unit_normals(self) -> Vector3D_Array:
         s_points, du_points, dv_points = self.get_surface_points_and_nudged_points()
         normals = np.cross(
             (du_points - s_points) / self.epsilon,

manim/mobject/opengl/opengl_three_dimensions.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any
+
 import numpy as np
 
 from manim.mobject.opengl.opengl_surface import OpenGLSurface
@@ -11,13 +13,13 @@
 class OpenGLSurfaceMesh(OpenGLVGroup):
     def __init__(
         self,
-        uv_surface,
-        resolution=None,
-        stroke_width=1,
-        normal_nudge=1e-2,
-        depth_test=True,
-        flat_stroke=False,
-        **kwargs,
+        uv_surface: OpenGLSurface,
+        resolution: tuple[int, int] | None = None,
+        stroke_width: float = 1,
+        normal_nudge: float = 1e-2,
+        depth_test: bool = True,
+        flat_stroke: bool = False,
+        **kwargs: Any,
     ):
         if not isinstance(uv_surface, OpenGLSurface):
             raise Exception("uv_surface must be of type OpenGLSurface")
@@ -31,7 +33,7 @@ def __init__(
             **kwargs,
         )
 
-    def init_points(self):
+    def init_points(self) -> None:
         uv_surface = self.uv_surface
 
         full_nu, full_nv = uv_surface.resolution

manim/mobject/opengl/opengl_vectorized_mobject.py

@@ -4,10 +4,11 @@
 import operator as op
 from collections.abc import Iterable, Sequence
 from functools import reduce, wraps
-from typing import Callable
+from typing import Any, Callable
 
 import moderngl
 import numpy as np
+from typing_extensions import Self
 
 from manim import config
 from manim.constants import *
@@ -171,6 +172,15 @@ def get_group_class(self):
     def get_mobject_type_class():
         return OpenGLVMobject
 
+    @property
+    def submobjects(self) -> Sequence[OpenGLVMobject]:
+        return self._submobjects if hasattr(self, "_submobjects") else []
+
+    @submobjects.setter
+    def submobjects(self, submobject_list: Iterable[OpenGLVMobject]) -> None:
+        self.remove(*self.submobjects)
+        self.add(*submobject_list)
+
     def init_data(self):
         super().init_data()
         self.data.pop("rgbas")
@@ -594,7 +604,9 @@ def set_points_as_corners(self, points: Iterable[float]) -> OpenGLVMobject:
         )
         return self
 
-    def set_points_smoothly(self, points, true_smooth=False):
+    def set_points_smoothly(
+        self, points: Point3DLike_Array, true_smooth: bool = False
+    ) -> Self:
         self.set_points_as_corners(points)
         self.make_smooth()
         return self
@@ -1654,7 +1666,7 @@ def construct(self):
                 self.add(circles_group)
     """
 
-    def __init__(self, *vmobjects, **kwargs):
+    def __init__(self, *vmobjects: OpenGLVMobject, **kwargs: Any):
         super().__init__(**kwargs)
         self.add(*vmobjects)
 

manim/renderer/vectorized_mobject_rendering.py

@@ -1,9 +1,18 @@
 from __future__ import annotations
 
 import collections
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from manim.renderer.opengl_renderer import (
+        OpenGLRenderer,
+        OpenGLVMobject,
+    )
+    from manim.typing import MatrixMN
+
 from ..utils import opengl
 from ..utils.space_ops import cross2d, earclip_triangulation
 from .shader import Shader
@@ -14,7 +23,9 @@
 ]
 
 
-def build_matrix_lists(mob):
+def build_matrix_lists(
+    mob: OpenGLVMobject,
+) -> collections.defaultdict[tuple[float, ...], list[OpenGLVMobject]]:
     root_hierarchical_matrix = mob.hierarchical_model_matrix()
     matrix_to_mobject_list = collections.defaultdict(list)
     if mob.has_points():
@@ -36,15 +47,21 @@ def build_matrix_lists(mob):
     return matrix_to_mobject_list
 
 
-def render_opengl_vectorized_mobject_fill(renderer, mobject):
+def render_opengl_vectorized_mobject_fill(
+    renderer: OpenGLRenderer, mobject: OpenGLVMobject
+) -> None:
     matrix_to_mobject_list = build_matrix_lists(mobject)
 
     for matrix_tuple, mobject_list in matrix_to_mobject_list.items():
         model_matrix = np.array(matrix_tuple).reshape((4, 4))
         render_mobject_fills_with_matrix(renderer, model_matrix, mobject_list)
 
 
-def render_mobject_fills_with_matrix(renderer, model_matrix, mobjects):
+def render_mobject_fills_with_matrix(
+    renderer: OpenGLRenderer,
+    model_matrix: MatrixMN,
+    mobjects: Iterable[OpenGLVMobject],
+) -> None:
     # Precompute the total number of vertices for which to reserve space.
     # Note that triangulate_mobject() will cache its results.
     total_size = 0
@@ -98,7 +115,7 @@ def render_mobject_fills_with_matrix(renderer, model_matrix, mobjects):
     vbo.release()
 
 
-def triangulate_mobject(mob):
+def triangulate_mobject(mob: OpenGLVMobject) -> np.ndarray:
     if not mob.needs_new_triangulation:
         return mob.triangulation
 
@@ -192,14 +209,20 @@ def triangulate_mobject(mob):
     return attributes
 
 
-def render_opengl_vectorized_mobject_stroke(renderer, mobject):
+def render_opengl_vectorized_mobject_stroke(
+    renderer: OpenGLRenderer, mobject: OpenGLVMobject
+) -> None:
     matrix_to_mobject_list = build_matrix_lists(mobject)
     for matrix_tuple, mobject_list in matrix_to_mobject_list.items():
         model_matrix = np.array(matrix_tuple).reshape((4, 4))
         render_mobject_strokes_with_matrix(renderer, model_matrix, mobject_list)
 
 
-def render_mobject_strokes_with_matrix(renderer, model_matrix, mobjects):
+def render_mobject_strokes_with_matrix(
+    renderer: OpenGLRenderer,
+    model_matrix: MatrixMN,
+    mobjects: Sequence[OpenGLVMobject],
+) -> None:
     # Precompute the total number of vertices for which to reserve space.
     total_size = 0
     for submob in mobjects:

mypy.ini

@@ -102,9 +102,6 @@ ignore_errors = True
 [mypy-manim.mobject.mobject]
 ignore_errors = True
 
-[mypy-manim.mobject.opengl.dot_cloud]
-ignore_errors = True
-
 [mypy-manim.mobject.opengl.opengl_compatibility]
 ignore_errors = True
 
@@ -123,9 +120,6 @@ ignore_errors = True
 [mypy-manim.mobject.opengl.opengl_surface]
 ignore_errors = True
 
-[mypy-manim.mobject.opengl.opengl_three_dimensions]
-ignore_errors = True
-
 [mypy-manim.mobject.opengl.opengl_vectorized_mobject]
 ignore_errors = True
 
@@ -162,9 +156,6 @@ ignore_errors = True
 [mypy-manim.renderer.shader_wrapper]
 ignore_errors = True
 
-[mypy-manim.renderer.vectorized_mobject_rendering]
-ignore_errors = True
-
 [mypy-manim.scene.three_d_scene]
 ignore_errors = True