Add type annotations to image_mobject.py (#4458)

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

manim/mobject/types/image_mobject.py

@@ -18,7 +18,13 @@ from ...camera.moving_camera import MovingCamera
 from ...constants import *
 from ...mobject.mobject import Mobject
 from ...utils.bezier import interpolate
-from ...utils.color import WHITE, ManimColor, color_to_int_rgb
+from ...utils.color import (
+    WHITE,
+    YELLOW_C,
+    ManimColor,
+    ParsableManimColor,
+    color_to_int_rgb,
+)
 from ...utils.images import change_to_rgba_array, get_full_raster_image_path
 
 __all__ = ["ImageMobject", "ImageMobjectFromCamera"]
@@ -62,9 +68,14 @@ class AbstractImageMobject(Mobject):
     def get_pixel_array(self) -> PixelArray:
         raise NotImplementedError()
 
-    def set_color(self, color, alpha=None, family=True):
+    def set_color(  # type: ignore[override]
+        self,
+        color: ParsableManimColor = YELLOW_C,
+        alpha: Any = None,
+        family: bool = True,
+    ) -> AbstractImageMobject:
         # Likely to be implemented in subclasses, but no obligation
-        pass
+        raise NotImplementedError()
 
     def set_resampling_algorithm(self, resampling_algorithm: int) -> Self:
         """
@@ -210,18 +221,23 @@ class ImageMobject(AbstractImageMobject):
         self.orig_alpha_pixel_array = self.pixel_array[:, :, 3].copy()
         super().__init__(scale_to_resolution, **kwargs)
 
-    def get_pixel_array(self):
+    def get_pixel_array(self) -> PixelArray:
         """A simple getter method."""
         return self.pixel_array
 
-    def set_color(self, color, alpha=None, family=True):
+    def set_color(  # type: ignore[override]
+        self,
+        color: ParsableManimColor = YELLOW_C,
+        alpha: Any = None,
+        family: bool = True,
+    ) -> Self:
         rgb = color_to_int_rgb(color)
         self.pixel_array[:, :, :3] = rgb
         if alpha is not None:
             self.pixel_array[:, :, 3] = int(255 * alpha)
         for submob in self.submobjects:
             submob.set_color(color, alpha, family)
-        self.color = color
+        self.color = ManimColor(color)
         return self
 
     def set_opacity(self, alpha: float) -> Self:
@@ -253,7 +269,7 @@ class ImageMobject(AbstractImageMobject):
         return self
 
     def interpolate_color(
-        self, mobject1: ImageMobject, mobject2: ImageMobject, alpha: float
+        self, mobject1: Mobject, mobject2: Mobject, alpha: float
     ) -> None:
         """Interpolates the array of pixel color values from one ImageMobject
         into an array of equal size in the target ImageMobject.
@@ -269,6 +285,8 @@ class ImageMobject(AbstractImageMobject):
         alpha
             Used to track the lerp relationship. Not opacity related.
         """
+        assert isinstance(mobject1, ImageMobject)
+        assert isinstance(mobject2, ImageMobject)
         assert mobject1.pixel_array.shape == mobject2.pixel_array.shape, (
             f"Mobject pixel array shapes incompatible for interpolation.\n"
             f"Mobject 1 ({mobject1}) : {mobject1.pixel_array.shape}\n"
@@ -292,7 +310,7 @@ class ImageMobject(AbstractImageMobject):
 
     def get_style(self) -> dict[str, Any]:
         return {
-            "fill_color": ManimColor(self.color.get_rgb()).to_hex(),
+            "fill_color": ManimColor(self.color.to_rgb()).to_hex(),
             "fill_opacity": self.fill_opacity,
         }
 
@@ -321,7 +339,7 @@ class ImageMobjectFromCamera(AbstractImageMobject):
         super().__init__(scale_to_resolution=False, **kwargs)
 
     # TODO: Get rid of this.
-    def get_pixel_array(self):
+    def get_pixel_array(self) -> PixelArray:
         self.pixel_array = self.camera.pixel_array
         return self.pixel_array
 
@@ -332,7 +350,11 @@ class ImageMobjectFromCamera(AbstractImageMobject):
         self.add(self.display_frame)
         return self
 
-    def interpolate_color(self, mobject1, mobject2, alpha) -> None:
+    def interpolate_color(
+        self, mobject1: Mobject, mobject2: Mobject, alpha: float
+    ) -> None:
+        assert isinstance(mobject1, ImageMobject)
+        assert isinstance(mobject2, ImageMobject)
         assert mobject1.pixel_array.shape == mobject2.pixel_array.shape, (
             f"Mobject pixel array shapes incompatible for interpolation.\n"
             f"Mobject 1 ({mobject1}) : {mobject1.pixel_array.shape}\n"

mypy.ini

@@ -103,9 +103,6 @@ ignore_errors = True
 [mypy-manim.mobject.table]
 ignore_errors = True
 
-[mypy-manim.mobject.types.image_mobject]
-ignore_errors = True
-
 [mypy-manim.mobject.types.point_cloud_mobject]
 ignore_errors = True