Add type annotations to mobject/graphing/number_line.py

Author: henrikmidtiby

manim/mobject/graphing/number_line.py

@@ -12,6 +12,10 @@
 from typing import TYPE_CHECKING, Callable
 
 if TYPE_CHECKING:
+    from typing import Any
+
+    from typing_extensions import Self
+
     from manim.mobject.geometry.tips import ArrowTip
     from manim.typing import Point3DLike
 
@@ -164,8 +168,8 @@ def __init__(
         decimal_number_config: dict | None = None,
         numbers_to_exclude: Iterable[float] | None = None,
         numbers_to_include: Iterable[float] | None = None,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> None:
         # avoid mutable arguments in defaults
         if numbers_to_exclude is None:
             numbers_to_exclude = []
@@ -189,6 +193,9 @@ def __init__(
 
         # turn into a NumPy array to scale by just applying the function
         self.x_range = np.array(x_range, dtype=float)
+        self.x_min: float
+        self.x_max: float
+        self.x_step: float
         self.x_min, self.x_max, self.x_step = scaling.function(self.x_range)
         self.length = length
         self.unit_size = unit_size
@@ -250,7 +257,9 @@ def __init__(
                     dict(
                         zip(
                             tick_range,
-                            self.scaling.get_custom_labels(
+                            # TODO:
+                            # Argument 2 to "zip" has incompatible type "Iterable[Mobject]"; expected "Iterable[str | float | VMobject]"  [arg-type]
+                            self.scaling.get_custom_labels(  # type: ignore[arg-type]
                                 tick_range,
                                 unit_decimal_places=decimal_number_config[
                                     "num_decimal_places"
@@ -267,21 +276,25 @@ def __init__(
                     font_size=self.font_size,
                 )
 
-    def rotate_about_zero(self, angle: float, axis: Sequence[float] = OUT, **kwargs):
+    def rotate_about_zero(
+        self, angle: float, axis: Sequence[float] = OUT, **kwargs: Any
+    ) -> Self:
         return self.rotate_about_number(0, angle, axis, **kwargs)
 
     def rotate_about_number(
-        self, number: float, angle: float, axis: Sequence[float] = OUT, **kwargs
-    ):
+        self, number: float, angle: float, axis: Sequence[float] = OUT, **kwargs: Any
+    ) -> Self:
         return self.rotate(angle, axis, about_point=self.n2p(number), **kwargs)
 
-    def add_ticks(self):
+    def add_ticks(self) -> None:
         """Adds ticks to the number line. Ticks can be accessed after creation
         via ``self.ticks``.
         """
         ticks = VGroup()
         elongated_tick_size = self.tick_size * self.longer_tick_multiple
-        elongated_tick_offsets = self.numbers_with_elongated_ticks - self.x_min
+        elongated_tick_offsets = (
+            np.array(self.numbers_with_elongated_ticks) - self.x_min
+        )
         for x in self.get_tick_range():
             size = self.tick_size
             if np.any(np.isclose(x - self.x_min, elongated_tick_offsets)):
@@ -413,19 +426,22 @@ def point_to_number(self, point: Sequence[float]) -> float:
         point = np.asarray(point)
         start, end = self.get_start_and_end()
         unit_vect = normalize(end - start)
-        proportion = np.dot(point - start, unit_vect) / np.dot(end - start, unit_vect)
+        proportion: float = np.dot(point - start, unit_vect) / np.dot(
+            end - start, unit_vect
+        )
         return interpolate(self.x_min, self.x_max, proportion)
 
     def n2p(self, number: float | np.ndarray) -> np.ndarray:
         """Abbreviation for :meth:`~.NumberLine.number_to_point`."""
         return self.number_to_point(number)
 
-    def p2n(self, point: Sequence[float]) -> float:
+    def p2n(self, point: Point3DLike) -> float:
         """Abbreviation for :meth:`~.NumberLine.point_to_number`."""
         return self.point_to_number(point)
 
     def get_unit_size(self) -> float:
-        return self.get_length() / (self.x_range[1] - self.x_range[0])
+        val: float = self.get_length() / (self.x_range[1] - self.x_range[0])
+        return val
 
     def get_unit_vector(self) -> np.ndarray:
         return super().get_unit_vector() * self.unit_size
@@ -436,8 +452,8 @@ def get_number_mobject(
         direction: Sequence[float] | None = None,
         buff: float | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-        **number_config,
+        label_constructor: type[MathTex] | None = None,
+        **number_config: dict[str, Any],
     ) -> VMobject:
         """Generates a positioned :class:`~.DecimalNumber` mobject
         generated according to ``label_constructor``.
@@ -476,7 +492,12 @@ def get_number_mobject(
             label_constructor = self.label_constructor
 
         num_mob = DecimalNumber(
-            x, font_size=font_size, mob_class=label_constructor, **number_config
+            # TODO:
+            # error: Argument 4 to "DecimalNumber" has incompatible type "**dict[str, dict[str, Any]]"; expected "int"  [arg-type]
+            x,
+            font_size=font_size,
+            mob_class=label_constructor,
+            **number_config,  # type: ignore[arg-type]
         )
 
         num_mob.next_to(self.number_to_point(x), direction=direction, buff=buff)
@@ -485,7 +506,7 @@ def get_number_mobject(
             num_mob.shift(num_mob[0].width * LEFT / 2)
         return num_mob
 
-    def get_number_mobjects(self, *numbers, **kwargs) -> VGroup:
+    def get_number_mobjects(self, *numbers: float, **kwargs: Any) -> VGroup:
         if len(numbers) == 0:
             numbers = self.default_numbers_to_display()
         return VGroup([self.get_number_mobject(number, **kwargs) for number in numbers])
@@ -498,9 +519,9 @@ def add_numbers(
         x_values: Iterable[float] | None = None,
         excluding: Iterable[float] | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-        **kwargs,
-    ):
+        label_constructor: type[MathTex] | None = None,
+        **kwargs: Any,
+    ) -> Self:
         """Adds :class:`~.DecimalNumber` mobjects representing their position
         at each tick of the number line. The numbers can be accessed after creation
         via ``self.numbers``.
@@ -551,11 +572,11 @@ def add_numbers(
     def add_labels(
         self,
         dict_values: dict[float, str | float | VMobject],
-        direction: Sequence[float] = None,
+        direction: Sequence[float] | None = None,
         buff: float | None = None,
         font_size: float | None = None,
-        label_constructor: VMobject | None = None,
-    ):
+        label_constructor: type[MathTex] | None = None,
+    ) -> Self:
         """Adds specifically positioned labels to the :class:`~.NumberLine` using a ``dict``.
         The labels can be accessed after creation via ``self.labels``.
 
@@ -598,6 +619,7 @@ def add_labels(
                 label = self._create_label_tex(label, label_constructor)
 
             if hasattr(label, "font_size"):
+                assert isinstance(label, MathTex)
                 label.font_size = font_size
             else:
                 raise AttributeError(f"{label} is not compatible with add_labels.")
@@ -612,7 +634,7 @@ def _create_label_tex(
         self,
         label_tex: str | float | VMobject,
         label_constructor: Callable | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> VMobject:
         """Checks if the label is a :class:`~.VMobject`, otherwise, creates a
         label by passing ``label_tex`` to ``label_constructor``.
@@ -633,24 +655,25 @@ def _create_label_tex(
         :class:`~.VMobject`
             The label.
         """
-        if label_constructor is None:
-            label_constructor = self.label_constructor
         if isinstance(label_tex, (VMobject, OpenGLVMobject)):
             return label_tex
-        else:
+        if label_constructor is None:
+            label_constructor = self.label_constructor
+        if isinstance(label_tex, str):
             return label_constructor(label_tex, **kwargs)
+        return label_constructor(str(label_tex), **kwargs)
 
     @staticmethod
-    def _decimal_places_from_step(step) -> int:
-        step = str(step)
-        if "." not in step:
+    def _decimal_places_from_step(step: float) -> int:
+        step_str = str(step)
+        if "." not in step_str:
             return 0
-        return len(step.split(".")[-1])
+        return len(step_str.split(".")[-1])
 
-    def __matmul__(self, other: float):
+    def __matmul__(self, other: float) -> np.ndarray:
         return self.n2p(other)
 
-    def __rmatmul__(self, other: Point3DLike | Mobject):
+    def __rmatmul__(self, other: Point3DLike | Mobject) -> float:
         if isinstance(other, Mobject):
             other = other.get_center()
         return self.p2n(other)
@@ -659,11 +682,11 @@ def __rmatmul__(self, other: Point3DLike | Mobject):
 class UnitInterval(NumberLine):
     def __init__(
         self,
-        unit_size=10,
-        numbers_with_elongated_ticks=None,
-        decimal_number_config=None,
-        **kwargs,
-    ):
+        unit_size: float = 10,
+        numbers_with_elongated_ticks: list[float] | None = None,
+        decimal_number_config: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
         numbers_with_elongated_ticks = (
             [0, 1]
             if numbers_with_elongated_ticks is None

manim/mobject/graphing/scale.py

@@ -11,7 +11,7 @@
 from manim.mobject.text.numbers import Integer
 
 if TYPE_CHECKING:
-    from typing import Callable
+    from typing import Callable, overload
 
     from manim.mobject.mobject import Mobject
 
@@ -28,6 +28,12 @@ class _ScaleBase:
     def __init__(self, custom_labels: bool = False):
         self.custom_labels = custom_labels
 
+    @overload
+    def function(self, value: float) -> float: ...
+
+    @overload
+    def function(self, value: np.array) -> np.array: ...
+
     def function(self, value: float) -> float:
         """The function that will be used to scale the values.
 
@@ -61,6 +67,7 @@ def inverse_function(self, value: float) -> float:
     def get_custom_labels(
         self,
         val_range: Iterable[float],
+        **kw_args: Any,
     ) -> Iterable[Mobject]:
         """Custom instructions for generating labels along an axis.
 

mypy.ini

@@ -73,6 +73,9 @@ ignore_errors = False
 [mypy-manim.mobject.graphing.scale.*]
 ignore_errors = False
 
+[mypy-manim.mobject.graphing.number_line.*]
+ignore_errors = False
+
 [mypy-manim.mobject.graphing.probability.*]
 ignore_errors = False