Add type annotations to transform_matching_parts.py (#4400)

* Add type annotations to `transform_matching_parts.py`

* Make two type errors quiet

* Make the pytests pass

---------

Co-authored-by: Francisco Manríquez Novoa <49853152+chopan050@users.noreply.github.com>

manim/animation/transform.py

@@ -835,7 +835,14 @@ class FadeTransform(Transform):
 
     """
 
-    def __init__(self, mobject, target_mobject, stretch=True, dim_to_match=1, **kwargs):
+    def __init__(
+        self,
+        mobject: Mobject,
+        target_mobject: Mobject,
+        stretch: bool = True,
+        dim_to_match: int = 1,
+        **kwargs: Any,
+    ):
         self.to_add_on_completion = target_mobject
         self.stretch = stretch
         self.dim_to_match = dim_to_match

manim/animation/transform_matching_parts.py

@@ -4,12 +4,13 @@ from __future__ import annotations
 
 __all__ = ["TransformMatchingShapes", "TransformMatchingTex"]
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from manim.mobject.opengl.opengl_mobject import OpenGLGroup, OpenGLMobject
 from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVGroup, OpenGLVMobject
+from manim.mobject.text.tex_mobject import MathTexPart
 
 from .._config import config
 from ..constants import RendererType
@@ -74,10 +75,10 @@ class TransformMatchingAbstractBase(AnimationGroup):
         transform_mismatches: bool = False,
         fade_transform_mismatches: bool = False,
         key_map: dict | None = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         if isinstance(mobject, OpenGLVMobject):
-            group_type = OpenGLVGroup
+            group_type: type[OpenGLVGroup | OpenGLGroup | VGroup | Group] = OpenGLVGroup
         elif isinstance(mobject, OpenGLMobject):
             group_type = OpenGLGroup
         elif isinstance(mobject, VMobject):
@@ -141,7 +142,7 @@ class TransformMatchingAbstractBase(AnimationGroup):
         self.to_add = target_mobject
 
     def get_shape_map(self, mobject: Mobject) -> dict:
-        shape_map = {}
+        shape_map: dict[int | str, VGroup | OpenGLVGroup] = {}
         for sm in self.get_mobject_parts(mobject):
             key = self.get_mobject_key(sm)
             if key not in shape_map:
@@ -149,23 +150,25 @@ class TransformMatchingAbstractBase(AnimationGroup):
                     shape_map[key] = OpenGLVGroup()
                 else:
                     shape_map[key] = VGroup()
-            shape_map[key].add(sm)
+            # error: Argument 1 to "add" of "OpenGLVGroup" has incompatible type "Mobject"; expected "OpenGLVMobject"  [arg-type]
+            shape_map[key].add(sm)  # type: ignore[arg-type]
         return shape_map
 
     def clean_up_from_scene(self, scene: Scene) -> None:
         # Interpolate all animations back to 0 to ensure source mobjects remain unchanged.
         for anim in self.animations:
             anim.interpolate(0)
-        scene.remove(self.mobject)
+        # error: Argument 1 to "remove" of "Scene" has incompatible type "OpenGLMobject"; expected "Mobject"  [arg-type]
+        scene.remove(self.mobject)  # type: ignore[arg-type]
         scene.remove(*self.to_remove)
         scene.add(self.to_add)
 
     @staticmethod
-    def get_mobject_parts(mobject: Mobject):
+    def get_mobject_parts(mobject: Mobject) -> list[Mobject]:
         raise NotImplementedError("To be implemented in subclass.")
 
     @staticmethod
-    def get_mobject_key(mobject: Mobject):
+    def get_mobject_key(mobject: Mobject) -> int | str:
         raise NotImplementedError("To be implemented in subclass.")
 
 
@@ -205,7 +208,7 @@ class TransformMatchingShapes(TransformMatchingAbstractBase):
         transform_mismatches: bool = False,
         fade_transform_mismatches: bool = False,
         key_map: dict | None = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(
             mobject,
@@ -269,7 +272,7 @@ class TransformMatchingTex(TransformMatchingAbstractBase):
         transform_mismatches: bool = False,
         fade_transform_mismatches: bool = False,
         key_map: dict | None = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(
             mobject,
@@ -294,4 +297,5 @@ class TransformMatchingTex(TransformMatchingAbstractBase):
 
     @staticmethod
     def get_mobject_key(mobject: Mobject) -> str:
+        assert isinstance(mobject, MathTexPart)
         return mobject.tex_string

mypy.ini

@@ -64,9 +64,6 @@ ignore_errors = True
 [mypy-manim.animation.speedmodifier]
 ignore_errors = True
 
-[mypy-manim.animation.transform_matching_parts]
-ignore_errors = True
-
 [mypy-manim.animation.transform]
 ignore_errors = True