Add type annotations to mobject.py

manim/camera/moving_camera.py

@@ -10,7 +10,7 @@
 __all__ = ["MovingCamera"]
 
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, Literal, overload
 
 from cairo import Context
 
@@ -20,7 +20,7 @@
 from ..camera.camera import Camera
 from ..constants import DOWN, LEFT, RIGHT, UP
 from ..mobject.frame import ScreenRectangle
-from ..mobject.mobject import Mobject
+from ..mobject.mobject import Mobject, _AnimationBuilder
 from ..utils.color import WHITE, ManimColor
 
 
@@ -166,13 +166,31 @@ def get_mobjects_indicating_movement(self) -> list[Mobject]:
         """
         return [self.frame]
 
+    @overload
+    def auto_zoom(
+        self,
+        mobjects: Iterable[Mobject],
+        margin: float,
+        only_mobjects_in_frame: bool,
+        animate: Literal[False],
+    ) -> Mobject: ...
+
+    @overload
+    def auto_zoom(
+        self,
+        mobjects: Iterable[Mobject],
+        margin: float,
+        only_mobjects_in_frame: bool,
+        animate: Literal[True],
+    ) -> _AnimationBuilder: ...
+
     def auto_zoom(
         self,
         mobjects: Iterable[Mobject],
         margin: float = 0,
         only_mobjects_in_frame: bool = False,
         animate: bool = True,
-    ) -> Mobject:
+    ) -> _AnimationBuilder | Mobject:
         """Zooms on to a given array of mobjects (or a singular mobject)
         and automatically resizes to frame all the mobjects.
 

manim/mobject/geometry/line.py

@@ -14,14 +14,14 @@
     "RightAngle",
 ]
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 
 from manim import config
 from manim.constants import *
 from manim.mobject.geometry.arc import Arc, ArcBetweenPoints, Dot, TipableVMobject
-from manim.mobject.geometry.tips import ArrowTriangleFilledTip
+from manim.mobject.geometry.tips import ArrowTip, ArrowTriangleFilledTip
 from manim.mobject.mobject import Mobject
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
@@ -648,9 +648,11 @@ def scale(self, factor: float, scale_tips: bool = False, **kwargs: Any) -> Self:
         self._set_stroke_width_from_length()
 
         if has_tip:
-            self.add_tip(tip=old_tips[0])
+            # error: Argument "tip" to "add_tip" of "TipableVMobject" has incompatible type "VMobject"; expected "ArrowTip | None"  [arg-type]
+            self.add_tip(tip=cast(ArrowTip, old_tips[0]))
         if has_start_tip:
-            self.add_tip(tip=old_tips[1], at_start=True)
+            # error: Argument "tip" to "add_tip" of "TipableVMobject" has incompatible type "VMobject"; expected "ArrowTip | None"  [arg-type]
+            self.add_tip(tip=cast(ArrowTip, old_tips[1]), at_start=True)
         return self
 
     def get_normal_vector(self) -> Vector3D:

manim/mobject/graphing/probability.py

@@ -207,13 +207,11 @@ def add_braces_and_labels(self) -> None:
                 if hasattr(parts, subattr):
                     self.add(getattr(parts, subattr))
 
-    def __getitem__(self, index: int) -> SampleSpace:
+    def __getitem__(self, index: int) -> VMobject:
         if hasattr(self, "horizontal_parts"):
-            val: SampleSpace = self.horizontal_parts[index]
-            return val
+            return self.horizontal_parts[index]
         elif hasattr(self, "vertical_parts"):
-            val = self.vertical_parts[index]
-            return val
+            return self.vertical_parts[index]
         return self.split()[index]
 
 

manim/mobject/matrix.py

@@ -40,15 +40,15 @@ def construct(self):
 
 
 import itertools as it
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable
 from typing import Any, Self
 
 import numpy as np
 
-from manim.mobject.mobject import Mobject
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
 from manim.mobject.text.numbers import DecimalNumber, Integer
 from manim.mobject.text.tex_mobject import MathTex, Tex
+from manim.typing import Vector2DLike, Vector3DLike
 
 from ..constants import *
 from ..mobject.types.vectorized_mobject import VGroup, VMobject
@@ -164,16 +164,16 @@ def construct(self):
 
     def __init__(
         self,
-        matrix: Iterable,
+        matrix: Iterable[Iterable[Any] | Vector2DLike],
         v_buff: float = 0.8,
         h_buff: float = 1.3,
         bracket_h_buff: float = MED_SMALL_BUFF,
         bracket_v_buff: float = MED_SMALL_BUFF,
         add_background_rectangles_to_entries: bool = False,
         include_background_rectangle: bool = False,
-        element_to_mobject: type[Mobject] | Callable[..., Mobject] = MathTex,
-        element_to_mobject_config: dict = {},
-        element_alignment_corner: Sequence[float] = DR,
+        element_to_mobject: type[VMobject] | Callable[..., VMobject] = MathTex,
+        element_to_mobject_config: dict[str, Any] = {},
+        element_alignment_corner: Vector3DLike = DR,
         left_bracket: str = "[",
         right_bracket: str = "]",
         stretch_brackets: bool = True,
@@ -206,7 +206,9 @@ def __init__(
         if self.include_background_rectangle:
             self.add_background_rectangle()
 
-    def _matrix_to_mob_matrix(self, matrix: np.ndarray) -> list[list[Mobject]]:
+    def _matrix_to_mob_matrix(
+        self, matrix: Iterable[Iterable[Any]]
+    ) -> list[list[VMobject]]:
         return [
             [
                 self.element_to_mobject(item, **self.element_to_mobject_config)
@@ -215,7 +217,7 @@ def _matrix_to_mob_matrix(self, matrix: np.ndarray) -> list[list[Mobject]]:
             for row in matrix
         ]
 
-    def _organize_mob_matrix(self, matrix: list[list[Mobject]]) -> Self:
+    def _organize_mob_matrix(self, matrix: list[list[VMobject]]) -> Self:
         for i, row in enumerate(matrix):
             for j, _ in enumerate(row):
                 mob = matrix[i][j]
@@ -401,7 +403,7 @@ def add_background_to_entries(self) -> Self:
             mob.add_background_rectangle()
         return self
 
-    def get_mob_matrix(self) -> list[list[Mobject]]:
+    def get_mob_matrix(self) -> list[list[VMobject]]:
         """Return the underlying mob matrix mobjects.
 
         Returns
@@ -483,8 +485,8 @@ def construct(self):
 
     def __init__(
         self,
-        matrix: Iterable,
-        element_to_mobject: type[Mobject] = DecimalNumber,
+        matrix: Iterable[Iterable[Any]],
+        element_to_mobject: type[VMobject] | Callable[..., VMobject] = DecimalNumber,
         element_to_mobject_config: dict[str, Any] = {"num_decimal_places": 1},
         **kwargs: Any,
     ):
@@ -528,8 +530,8 @@ def construct(self):
 
     def __init__(
         self,
-        matrix: Iterable,
-        element_to_mobject: type[Mobject] = Integer,
+        matrix: Iterable[Iterable[Any]],
+        element_to_mobject: type[VMobject] | Callable[..., VMobject] = Integer,
         **kwargs: Any,
     ):
         """
@@ -566,8 +568,8 @@ def construct(self):
 
     def __init__(
         self,
-        matrix: Iterable,
-        element_to_mobject: type[Mobject] | Callable[..., Mobject] = lambda m: m,
+        matrix: Iterable[Iterable[Any]],
+        element_to_mobject: type[VMobject] | Callable[..., VMobject] = lambda m: m,
         **kwargs: Any,
     ):
         super().__init__(matrix, element_to_mobject=element_to_mobject, **kwargs)