diff --git a/manimlib/mobject/types/dot_cloud.py b/manimlib/mobject/types/dot_cloud.py index 83aa9b07..ad65647c 100644 --- a/manimlib/mobject/types/dot_cloud.py +++ b/manimlib/mobject/types/dot_cloud.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: import numpy.typing as npt - from typing import Sequence, Tuple + from typing import Sequence, Tuple, Self from manimlib.typing import ManimColor, Vect3, Vect3Array @@ -70,7 +70,7 @@ class DotCloud(PMobject): v_buff_ratio: float = 1.0, d_buff_ratio: float = 1.0, height: float = DEFAULT_GRID_HEIGHT, - ): + ) -> Self: n_points = n_rows * n_cols * n_layers points = np.repeat(range(n_points), 3, axis=0).reshape((n_points, 3)) points[:, 0] = points[:, 0] % n_cols @@ -96,7 +96,7 @@ class DotCloud(PMobject): return self @Mobject.affects_data - def set_radii(self, radii: npt.ArrayLike): + def set_radii(self, radii: npt.ArrayLike) -> Self: n_points = self.get_num_points() radii = np.array(radii).reshape((len(radii), 1)) self.data["radius"][:] = resize_with_interpolation(radii, n_points) @@ -107,7 +107,7 @@ class DotCloud(PMobject): return self.data["radius"] @Mobject.affects_data - def set_radius(self, radius: float): + def set_radius(self, radius: float) -> Self: data = self.data if self.get_num_points() > 0 else self._data_defaults data["radius"][:] = radius self.refresh_bounding_box() @@ -116,13 +116,14 @@ class DotCloud(PMobject): def get_radius(self) -> float: return self.get_radii().max() - def set_glow_factor(self, glow_factor: float) -> None: + def set_glow_factor(self, glow_factor: float) -> Self: self.uniforms["glow_factor"] = glow_factor + return self def get_glow_factor(self) -> float: return self.uniforms["glow_factor"] - def compute_bounding_box(self) -> np.ndarray: + def compute_bounding_box(self) -> Vect3Array: bb = super().compute_bounding_box() radius = self.get_radius() bb[0] += np.full((3,), -radius) @@ -134,7 +135,7 @@ class DotCloud(PMobject): scale_factor: float | npt.ArrayLike, scale_radii: bool = True, **kwargs - ): + ) -> Self: super().scale(scale_factor, **kwargs) if scale_radii: self.set_radii(scale_factor * self.get_radii()) @@ -145,7 +146,7 @@ class DotCloud(PMobject): reflectiveness: float = 0.5, gloss: float = 0.1, shadow: float = 0.2 - ): + ) -> Self: self.set_shading(reflectiveness, gloss, shadow) self.apply_depth_test() return self diff --git a/manimlib/mobject/types/point_cloud_mobject.py b/manimlib/mobject/types/point_cloud_mobject.py index f6180b3f..ab958473 100644 --- a/manimlib/mobject/types/point_cloud_mobject.py +++ b/manimlib/mobject/types/point_cloud_mobject.py @@ -10,7 +10,7 @@ from manimlib.utils.iterables import resize_with_interpolation from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Callable + from typing import Callable, Self from manimlib.typing import ManimColor, Vect3, Vect3Array, Vect4Array @@ -28,7 +28,7 @@ class PMobject(Mobject): rgbas: Vect4Array | None = None, color: ManimColor | None = None, opacity: float | None = None - ): + ) -> Self: """ points must be a Nx3 numpy array, as must rgbas if it is not None """ @@ -46,13 +46,13 @@ class PMobject(Mobject): self.data["rgba"][-len(rgbas):] = rgbas return self - def add_point(self, point: Vect3, rgba=None, color=None, opacity=None): + def add_point(self, point: Vect3, rgba=None, color=None, opacity=None) -> Self: rgbas = None if rgba is None else [rgba] self.add_points([point], rgbas, color, opacity) return self @Mobject.affects_data - def set_color_by_gradient(self, *colors: ManimColor): + def set_color_by_gradient(self, *colors: ManimColor) -> Self: self.data["rgba"][:] = np.array(list(map( color_to_rgba, color_gradient(colors, self.get_num_points()) @@ -60,20 +60,20 @@ class PMobject(Mobject): return self @Mobject.affects_data - def match_colors(self, pmobject: PMobject): + def match_colors(self, pmobject: PMobject) -> Self: self.data["rgba"][:] = resize_with_interpolation( pmobject.data["rgba"], self.get_num_points() ) return self @Mobject.affects_data - def filter_out(self, condition: Callable[[np.ndarray], bool]): + def filter_out(self, condition: Callable[[np.ndarray], bool]) -> Self: for mob in self.family_members_with_points(): mob.data = mob.data[~np.apply_along_axis(condition, 1, mob.get_points())] return self @Mobject.affects_data - def sort_points(self, function: Callable[[Vect3], None] = lambda p: p[0]): + def sort_points(self, function: Callable[[Vect3], None] = lambda p: p[0]) -> Self: """ function is any map from R^3 to R """ @@ -85,7 +85,7 @@ class PMobject(Mobject): return self @Mobject.affects_data - def ingest_submobjects(self): + def ingest_submobjects(self) -> Self: self.data = np.vstack([ sm.data for sm in self.get_family() ]) @@ -96,7 +96,7 @@ class PMobject(Mobject): return self.get_points()[int(index)] @Mobject.affects_data - def pointwise_become_partial(self, pmobject: PMobject, a: float, b: float): + def pointwise_become_partial(self, pmobject: PMobject, a: float, b: float) -> Self: lower_index = int(a * pmobject.get_num_points()) upper_index = int(b * pmobject.get_num_points()) self.data = pmobject.data[lower_index:upper_index].copy()