manim 0.18.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

manim/mobject/types/image_mobject.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 __all__ = ["AbstractImageMobject", "ImageMobject", "ImageMobjectFromCamera"]
 
 import pathlib
+from typing import TYPE_CHECKING
 
 import numpy as np
 from PIL import Image
@@ -21,6 +22,14 @@ from ...utils.images import change_to_rgba_array, get_full_raster_image_path
 
 __all__ = ["ImageMobject", "ImageMobjectFromCamera"]
 
+if TYPE_CHECKING:
+    from typing import Any
+
+    import numpy.typing as npt
+    from typing_extensions import Self
+
+    from manim.typing import StrPath
+
 
 class AbstractImageMobject(Mobject):
     """
@@ -39,23 +48,23 @@ class AbstractImageMobject(Mobject):
     def __init__(
         self,
         scale_to_resolution: int,
-        pixel_array_dtype="uint8",
-        resampling_algorithm=Resampling.BICUBIC,
-        **kwargs,
-    ):
+        pixel_array_dtype: str = "uint8",
+        resampling_algorithm: Resampling = Resampling.BICUBIC,
+        **kwargs: Any,
+    ) -> None:
         self.pixel_array_dtype = pixel_array_dtype
         self.scale_to_resolution = scale_to_resolution
         self.set_resampling_algorithm(resampling_algorithm)
         super().__init__(**kwargs)
 
-    def get_pixel_array(self):
+    def get_pixel_array(self) -> None:
         raise NotImplementedError()
 
     def set_color(self, color, alpha=None, family=True):
         # Likely to be implemented in subclasses, but no obligation
         pass
 
-    def set_resampling_algorithm(self, resampling_algorithm: int):
+    def set_resampling_algorithm(self, resampling_algorithm: int) -> Self:
         """
         Sets the interpolation method for upscaling the image. By default the image is
         interpolated using bicubic algorithm. This method lets you change it.
@@ -87,7 +96,7 @@ class AbstractImageMobject(Mobject):
             )
         return self
 
-    def reset_points(self):
+    def reset_points(self) -> None:
         """Sets :attr:`points` to be the four image corners."""
         self.points = np.array(
             [
@@ -171,15 +180,15 @@ class ImageMobject(AbstractImageMobject):
 
     def __init__(
         self,
-        filename_or_array,
+        filename_or_array: StrPath | npt.NDArray,
         scale_to_resolution: int = QUALITIES[DEFAULT_QUALITY]["pixel_height"],
-        invert=False,
-        image_mode="RGBA",
-        **kwargs,
-    ):
-        self.fill_opacity = 1
-        self.stroke_opacity = 1
-        self.invert = invert
+        invert: bool = False,
+        image_mode: str = "RGBA",
+        **kwargs: Any,
+    ) -> None:
+        self.fill_opacity: float = 1
+        self.stroke_opacity: float = 1
+        self.invert_image = invert
         self.image_mode = image_mode
         if isinstance(filename_or_array, (str, pathlib.PurePath)):
             path = get_full_raster_image_path(filename_or_array)
@@ -192,7 +201,7 @@ class ImageMobject(AbstractImageMobject):
         self.pixel_array = change_to_rgba_array(
             self.pixel_array, self.pixel_array_dtype
         )
-        if self.invert:
+        if self.invert_image:
             self.pixel_array[:, :, :3] = (
                 np.iinfo(self.pixel_array_dtype).max - self.pixel_array[:, :, :3]
             )
@@ -212,7 +221,7 @@ class ImageMobject(AbstractImageMobject):
         self.color = color
         return self
 
-    def set_opacity(self, alpha: float):
+    def set_opacity(self, alpha: float) -> Self:
         """Sets the image's opacity.
 
         Parameters
@@ -226,7 +235,7 @@ class ImageMobject(AbstractImageMobject):
         self.stroke_opacity = alpha
         return self
 
-    def fade(self, darkness: float = 0.5, family: bool = True):
+    def fade(self, darkness: float = 0.5, family: bool = True) -> Self:
         """Sets the image's opacity using a 1 - alpha relationship.
 
         Parameters
@@ -243,7 +252,7 @@ class ImageMobject(AbstractImageMobject):
 
     def interpolate_color(
         self, mobject1: ImageMobject, mobject2: ImageMobject, alpha: float
-    ):
+    ) -> None:
         """Interpolates the array of pixel color values from one ImageMobject
         into an array of equal size in the target ImageMobject.
 
@@ -279,7 +288,7 @@ class ImageMobject(AbstractImageMobject):
             alpha,
         ).astype(self.pixel_array_dtype)
 
-    def get_style(self):
+    def get_style(self) -> dict[str, Any]:
         return {
             "fill_color": ManimColor(self.color.get_rgb()).to_hex(),
             "fill_opacity": self.fill_opacity,
@@ -292,7 +301,12 @@ class ImageMobject(AbstractImageMobject):
 
 
 class ImageMobjectFromCamera(AbstractImageMobject):
-    def __init__(self, camera, default_display_frame_config=None, **kwargs):
+    def __init__(
+        self,
+        camera,
+        default_display_frame_config: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
         self.camera = camera
         if default_display_frame_config is None:
             default_display_frame_config = {
@@ -309,14 +323,14 @@ class ImageMobjectFromCamera(AbstractImageMobject):
         self.pixel_array = self.camera.pixel_array
         return self.pixel_array
 
-    def add_display_frame(self, **kwargs):
+    def add_display_frame(self, **kwargs: Any) -> Self:
         config = dict(self.default_display_frame_config)
         config.update(kwargs)
         self.display_frame = SurroundingRectangle(self, **config)
         self.add(self.display_frame)
         return self
 
-    def interpolate_color(self, mobject1, mobject2, alpha):
+    def interpolate_color(self, mobject1, mobject2, alpha) -> None:
         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"

manim/mobject/types/point_cloud_mobject.py

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 __all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
@@ -17,6 +19,7 @@ from ...utils.color import (
     WHITE,
     YELLOW,
     ManimColor,
+    ParsableManimColor,
     color_gradient,
     color_to_rgba,
     rgba_to_color,
@@ -25,6 +28,15 @@ from ...utils.iterables import stretch_array_to_length
 
 __all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    import numpy.typing as npt
+    from typing_extensions import Self
+
+    from manim.typing import ManimFloat, Point3DLike, Vector3D
+
 
 class PMobject(Mobject, metaclass=ConvertToOpenGL):
     """A disc made of a cloud of Dots
@@ -55,19 +67,25 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
     """
 
-    def __init__(self, stroke_width=DEFAULT_STROKE_WIDTH, **kwargs):
+    def __init__(self, stroke_width: int = DEFAULT_STROKE_WIDTH, **kwargs: Any) -> None:
         self.stroke_width = stroke_width
         super().__init__(**kwargs)
 
-    def reset_points(self):
+    def reset_points(self) -> Self:
         self.rgbas = np.zeros((0, 4))
         self.points = np.zeros((0, 3))
         return self
 
-    def get_array_attrs(self):
+    def get_array_attrs(self) -> list[str]:
         return super().get_array_attrs() + ["rgbas"]
 
-    def add_points(self, points, rgbas=None, color=None, alpha=1):
+    def add_points(
+        self,
+        points: npt.NDArray,
+        rgbas: npt.NDArray | None = None,
+        color: ParsableManimColor | None = None,
+        alpha: float = 1,
+    ) -> Self:
         """Add points.
 
         Points must be a Nx3 numpy array.
@@ -85,24 +103,26 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         self.rgbas = np.append(self.rgbas, rgbas, axis=0)
         return self
 
-    def set_color(self, color=YELLOW, family=True):
+    def set_color(
+        self, color: ParsableManimColor = YELLOW, family: bool = True
+    ) -> Self:
         rgba = color_to_rgba(color)
         mobs = self.family_members_with_points() if family else [self]
         for mob in mobs:
             mob.rgbas[:, :] = rgba
-        self.color = color
+        self.color = ManimColor.parse(color)
         return self
 
-    def get_stroke_width(self):
+    def get_stroke_width(self) -> int:
         return self.stroke_width
 
-    def set_stroke_width(self, width, family=True):
+    def set_stroke_width(self, width: int, family: bool = True) -> Self:
         mobs = self.family_members_with_points() if family else [self]
         for mob in mobs:
             mob.stroke_width = width
         return self
 
-    def set_color_by_gradient(self, *colors):
+    def set_color_by_gradient(self, *colors: ParsableManimColor) -> Self:
         self.rgbas = np.array(
             list(map(color_to_rgba, color_gradient(*colors, len(self.points)))),
         )
@@ -110,11 +130,11 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
     def set_colors_by_radial_gradient(
         self,
-        center=None,
-        radius=1,
-        inner_color=WHITE,
-        outer_color=BLACK,
-    ):
+        center: Point3DLike | None = None,
+        radius: float = 1,
+        inner_color: ParsableManimColor = WHITE,
+        outer_color: ParsableManimColor = BLACK,
+    ) -> Self:
         start_rgba, end_rgba = list(map(color_to_rgba, [inner_color, outer_color]))
         if center is None:
             center = self.get_center()
@@ -129,48 +149,48 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
             )
         return self
 
-    def match_colors(self, mobject):
+    def match_colors(self, mobject: Mobject) -> Self:
         Mobject.align_data(self, mobject)
         self.rgbas = np.array(mobject.rgbas)
         return self
 
-    def filter_out(self, condition):
+    def filter_out(self, condition: npt.NDArray) -> Self:
         for mob in self.family_members_with_points():
             to_eliminate = ~np.apply_along_axis(condition, 1, mob.points)
             mob.points = mob.points[to_eliminate]
             mob.rgbas = mob.rgbas[to_eliminate]
         return self
 
-    def thin_out(self, factor=5):
-        """
-        Removes all but every nth point for n = factor
-        """
+    def thin_out(self, factor: int = 5) -> Self:
+        """Removes all but every nth point for n = factor"""
         for mob in self.family_members_with_points():
             num_points = self.get_num_points()
             mob.apply_over_attr_arrays(
-                lambda arr: arr[np.arange(0, num_points, factor)],
+                lambda arr, n=num_points: arr[np.arange(0, n, factor)],
             )
         return self
 
-    def sort_points(self, function=lambda p: p[0]):
-        """
-        Function is any map from R^3 to R
-        """
+    def sort_points(
+        self, function: Callable[[npt.NDArray[ManimFloat]], float] = lambda p: p[0]
+    ) -> Self:
+        """Function is any map from R^3 to R"""
         for mob in self.family_members_with_points():
             indices = np.argsort(np.apply_along_axis(function, 1, mob.points))
-            mob.apply_over_attr_arrays(lambda arr: arr[indices])
+            mob.apply_over_attr_arrays(lambda arr, idx=indices: arr[idx])
         return self
 
-    def fade_to(self, color, alpha, family=True):
+    def fade_to(
+        self, color: ParsableManimColor, alpha: float, family: bool = True
+    ) -> Self:
         self.rgbas = interpolate(self.rgbas, color_to_rgba(color), alpha)
         for mob in self.submobjects:
             mob.fade_to(color, alpha, family)
         return self
 
-    def get_all_rgbas(self):
+    def get_all_rgbas(self) -> npt.NDArray:
         return self.get_merged_array("rgbas")
 
-    def ingest_submobjects(self):
+    def ingest_submobjects(self) -> Self:
         attrs = self.get_array_attrs()
         arrays = list(map(self.get_merged_array, attrs))
         for attr, array in zip(attrs, arrays):
@@ -178,30 +198,32 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         self.submobjects = []
         return self
 
-    def get_color(self):
+    def get_color(self) -> ManimColor:
         return rgba_to_color(self.rgbas[0, :])
 
-    def point_from_proportion(self, alpha):
+    def point_from_proportion(self, alpha: float) -> Any:
         index = alpha * (self.get_num_points() - 1)
-        return self.points[index]
+        return self.points[np.floor(index)]
 
     @staticmethod
-    def get_mobject_type_class():
+    def get_mobject_type_class() -> type[PMobject]:
         return PMobject
 
     # Alignment
-    def align_points_with_larger(self, larger_mobject):
+    def align_points_with_larger(self, larger_mobject: Mobject) -> None:
         assert isinstance(larger_mobject, PMobject)
         self.apply_over_attr_arrays(
             lambda a: stretch_array_to_length(a, larger_mobject.get_num_points()),
         )
 
-    def get_point_mobject(self, center=None):
+    def get_point_mobject(self, center: Point3DLike | None = None) -> Point:
         if center is None:
             center = self.get_center()
         return Point(center)
 
-    def interpolate_color(self, mobject1, mobject2, alpha):
+    def interpolate_color(
+        self, mobject1: Mobject, mobject2: Mobject, alpha: float
+    ) -> Self:
         self.rgbas = interpolate(mobject1.rgbas, mobject2.rgbas, alpha)
         self.set_stroke_width(
             interpolate(
@@ -212,7 +234,7 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
         )
         return self
 
-    def pointwise_become_partial(self, mobject, a, b):
+    def pointwise_become_partial(self, mobject: Mobject, a: float, b: float) -> None:
         lower_index, upper_index = (int(x * mobject.get_num_points()) for x in (a, b))
         for attr in self.get_array_attrs():
             full_array = getattr(mobject, attr)
@@ -222,24 +244,31 @@ class PMobject(Mobject, metaclass=ConvertToOpenGL):
 
 # TODO, Make the two implementations below non-redundant
 class Mobject1D(PMobject, metaclass=ConvertToOpenGL):
-    def __init__(self, density=DEFAULT_POINT_DENSITY_1D, **kwargs):
+    def __init__(self, density: int = DEFAULT_POINT_DENSITY_1D, **kwargs: Any) -> None:
         self.density = density
         self.epsilon = 1.0 / self.density
         super().__init__(**kwargs)
 
-    def add_line(self, start, end, color=None):
+    def add_line(
+        self,
+        start: npt.NDArray,
+        end: npt.NDArray,
+        color: ParsableManimColor | None = None,
+    ) -> None:
         start, end = list(map(np.array, [start, end]))
         length = np.linalg.norm(end - start)
         if length == 0:
-            points = [start]
+            points = np.array([start])
         else:
             epsilon = self.epsilon / length
-            points = [interpolate(start, end, t) for t in np.arange(0, 1, epsilon)]
+            points = np.array(
+                [interpolate(start, end, t) for t in np.arange(0, 1, epsilon)]
+            )
         self.add_points(points, color=color)
 
 
 class Mobject2D(PMobject, metaclass=ConvertToOpenGL):
-    def __init__(self, density=DEFAULT_POINT_DENSITY_2D, **kwargs):
+    def __init__(self, density: int = DEFAULT_POINT_DENSITY_2D, **kwargs: Any) -> None:
         self.density = density
         self.epsilon = 1.0 / self.density
         super().__init__(**kwargs)
@@ -268,7 +297,7 @@ class PGroup(PMobject):
 
     """
 
-    def __init__(self, *pmobs, **kwargs):
+    def __init__(self, *pmobs: Any, **kwargs: Any) -> None:
         if not all(isinstance(m, (PMobject, OpenGLPMobject)) for m in pmobs):
             raise ValueError(
                 "All submobjects must be of type PMobject or OpenGLPMObject"
@@ -277,10 +306,13 @@ class PGroup(PMobject):
         super().__init__(**kwargs)
         self.add(*pmobs)
 
-    def fade_to(self, color, alpha, family=True):
+    def fade_to(
+        self, color: ParsableManimColor, alpha: float, family: bool = True
+    ) -> Self:
         if family:
             for mob in self.submobjects:
                 mob.fade_to(color, alpha, family)
+        return self
 
 
 class PointCloudDot(Mobject1D):
@@ -317,13 +349,13 @@ class PointCloudDot(Mobject1D):
 
     def __init__(
         self,
-        center=ORIGIN,
-        radius=2.0,
-        stroke_width=2,
-        density=DEFAULT_POINT_DENSITY_1D,
-        color=YELLOW,
-        **kwargs,
-    ):
+        center: Vector3D = ORIGIN,
+        radius: float = 2.0,
+        stroke_width: int = 2,
+        density: int = DEFAULT_POINT_DENSITY_1D,
+        color: ManimColor = YELLOW,
+        **kwargs: Any,
+    ) -> None:
         self.radius = radius
         self.epsilon = 1.0 / density
         super().__init__(
@@ -331,22 +363,24 @@ class PointCloudDot(Mobject1D):
         )
         self.shift(center)
 
-    def init_points(self):
+    def init_points(self) -> None:
         self.reset_points()
         self.generate_points()
 
-    def generate_points(self):
+    def generate_points(self) -> None:
         self.add_points(
-            [
-                r * (np.cos(theta) * RIGHT + np.sin(theta) * UP)
-                for r in np.arange(self.epsilon, self.radius, self.epsilon)
-                # Num is equal to int(stop - start)/ (step + 1) reformulated.
-                for theta in np.linspace(
-                    0,
-                    2 * np.pi,
-                    num=int(2 * np.pi * (r + self.epsilon) / self.epsilon),
-                )
-            ],
+            np.array(
+                [
+                    r * (np.cos(theta) * RIGHT + np.sin(theta) * UP)
+                    for r in np.arange(self.epsilon, self.radius, self.epsilon)
+                    # Num is equal to int(stop - start)/ (step + 1) reformulated.
+                    for theta in np.linspace(
+                        0,
+                        2 * np.pi,
+                        num=int(2 * np.pi * (r + self.epsilon) / self.epsilon),
+                    )
+                ]
+            ),
         )
 
 
@@ -371,14 +405,16 @@ class Point(PMobject):
                 self.add(point)
     """
 
-    def __init__(self, location=ORIGIN, color=BLACK, **kwargs):
+    def __init__(
+        self, location: Vector3D = ORIGIN, color: ManimColor = BLACK, **kwargs: Any
+    ) -> None:
         self.location = location
         super().__init__(color=color, **kwargs)
 
-    def init_points(self):
+    def init_points(self) -> None:
         self.reset_points()
         self.generate_points()
         self.set_points([self.location])
 
-    def generate_points(self):
-        self.add_points([self.location])
+    def generate_points(self) -> None:
+        self.add_points(np.array([self.location]))