manim 0.17.0__py3-none-any.whl → 0.19.1__py3-none-any.whl

manim/mobject/graphing/probability.py

@@ -5,20 +5,20 @@ from __future__ import annotations
 __all__ = ["SampleSpace", "BarChart"]
 
 
-from typing import Iterable, MutableSequence, Sequence
+from collections.abc import Iterable, MutableSequence, Sequence
+from typing import Any
 
 import numpy as np
-from colour import Color
 
 from manim import config, logger
 from manim.constants import *
 from manim.mobject.geometry.polygram import Rectangle
 from manim.mobject.graphing.coordinate_systems import Axes
-from manim.mobject.mobject import Mobject
-from manim.mobject.opengl.opengl_mobject import OpenGLMobject
+from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
 from manim.mobject.svg.brace import Brace
 from manim.mobject.text.tex_mobject import MathTex, Tex
 from manim.mobject.types.vectorized_mobject import VGroup, VMobject
+from manim.typing import Vector3D
 from manim.utils.color import (
     BLUE_E,
     DARK_GREY,
@@ -26,6 +26,7 @@ from manim.utils.color import (
     LIGHT_GREY,
     MAROON_B,
     YELLOW,
+    ParsableManimColor,
     color_gradient,
 )
 from manim.utils.iterables import tuplify
@@ -34,7 +35,8 @@ EPSILON = 0.0001
 
 
 class SampleSpace(Rectangle):
-    """
+    """A mobject representing a twodimensional rectangular
+    sampling space.
 
     Examples
     --------
@@ -53,13 +55,13 @@ class SampleSpace(Rectangle):
 
     def __init__(
         self,
-        height=3,
-        width=3,
-        fill_color=DARK_GREY,
-        fill_opacity=1,
-        stroke_width=0.5,
-        stroke_color=LIGHT_GREY,
-        default_label_scale_val=1,
+        height: float = 3,
+        width: float = 3,
+        fill_color: ParsableManimColor = DARK_GREY,
+        fill_opacity: float = 1,
+        stroke_width: float = 0.5,
+        stroke_color: ParsableManimColor = LIGHT_GREY,
+        default_label_scale_val: float = 1,
     ):
         super().__init__(
             height=height,
@@ -71,7 +73,9 @@ class SampleSpace(Rectangle):
         )
         self.default_label_scale_val = default_label_scale_val
 
-    def add_title(self, title="Sample space", buff=MED_SMALL_BUFF):
+    def add_title(
+        self, title: str = "Sample space", buff: float = MED_SMALL_BUFF
+    ) -> None:
         # TODO, should this really exist in SampleSpaceScene
         title_mob = Tex(title)
         if title_mob.width > self.width:
@@ -80,23 +84,30 @@ class SampleSpace(Rectangle):
         self.title = title_mob
         self.add(title_mob)
 
-    def add_label(self, label):
+    def add_label(self, label: str) -> None:
         self.label = label
 
-    def complete_p_list(self, p_list):
-        new_p_list = list(tuplify(p_list))
+    def complete_p_list(self, p_list: float | Iterable[float]) -> list[float]:
+        p_list_tuplified: tuple[float] = tuplify(p_list)
+        new_p_list = list(p_list_tuplified)
         remainder = 1.0 - sum(new_p_list)
         if abs(remainder) > EPSILON:
             new_p_list.append(remainder)
         return new_p_list
 
-    def get_division_along_dimension(self, p_list, dim, colors, vect):
-        p_list = self.complete_p_list(p_list)
-        colors = color_gradient(colors, len(p_list))
+    def get_division_along_dimension(
+        self,
+        p_list: float | Iterable[float],
+        dim: int,
+        colors: Sequence[ParsableManimColor],
+        vect: Vector3D,
+    ) -> VGroup:
+        p_list_complete = self.complete_p_list(p_list)
+        colors_in_gradient = color_gradient(colors, len(p_list_complete))
 
         last_point = self.get_edge_center(-vect)
         parts = VGroup()
-        for factor, color in zip(p_list, colors):
+        for factor, color in zip(p_list_complete, colors_in_gradient):
             part = SampleSpace()
             part.set_fill(color, 1)
             part.replace(self, stretch=True)
@@ -106,33 +117,43 @@ class SampleSpace(Rectangle):
             parts.add(part)
         return parts
 
-    def get_horizontal_division(self, p_list, colors=[GREEN_E, BLUE_E], vect=DOWN):
+    def get_horizontal_division(
+        self,
+        p_list: float | Iterable[float],
+        colors: Sequence[ParsableManimColor] = [GREEN_E, BLUE_E],
+        vect: Vector3D = DOWN,
+    ) -> VGroup:
         return self.get_division_along_dimension(p_list, 1, colors, vect)
 
-    def get_vertical_division(self, p_list, colors=[MAROON_B, YELLOW], vect=RIGHT):
+    def get_vertical_division(
+        self,
+        p_list: float | Iterable[float],
+        colors: Sequence[ParsableManimColor] = [MAROON_B, YELLOW],
+        vect: Vector3D = RIGHT,
+    ) -> VGroup:
         return self.get_division_along_dimension(p_list, 0, colors, vect)
 
-    def divide_horizontally(self, *args, **kwargs):
+    def divide_horizontally(self, *args: Any, **kwargs: Any) -> None:
         self.horizontal_parts = self.get_horizontal_division(*args, **kwargs)
         self.add(self.horizontal_parts)
 
-    def divide_vertically(self, *args, **kwargs):
+    def divide_vertically(self, *args: Any, **kwargs: Any) -> None:
         self.vertical_parts = self.get_vertical_division(*args, **kwargs)
         self.add(self.vertical_parts)
 
     def get_subdivision_braces_and_labels(
         self,
-        parts,
-        labels,
-        direction,
-        buff=SMALL_BUFF,
-        min_num_quads=1,
-    ):
+        parts: VGroup,
+        labels: list[str | VMobject | OpenGLVMobject],
+        direction: Vector3D,
+        buff: float = SMALL_BUFF,
+        min_num_quads: int = 1,
+    ) -> VGroup:
         label_mobs = VGroup()
         braces = VGroup()
         for label, part in zip(labels, parts):
             brace = Brace(part, direction, min_num_quads=min_num_quads, buff=buff)
-            if isinstance(label, (Mobject, OpenGLMobject)):
+            if isinstance(label, (VMobject, OpenGLVMobject)):
                 label_mob = label
             else:
                 label_mob = MathTex(label)
@@ -140,34 +161,44 @@ class SampleSpace(Rectangle):
             label_mob.next_to(brace, direction, buff)
 
             braces.add(brace)
+            assert isinstance(label_mob, VMobject)
             label_mobs.add(label_mob)
-        parts.braces = braces
-        parts.labels = label_mobs
-        parts.label_kwargs = {
+        parts.braces = braces  # type: ignore[attr-defined]
+        parts.labels = label_mobs  # type: ignore[attr-defined]
+        parts.label_kwargs = {  # type: ignore[attr-defined]
             "labels": label_mobs.copy(),
             "direction": direction,
             "buff": buff,
         }
         return VGroup(parts.braces, parts.labels)
 
-    def get_side_braces_and_labels(self, labels, direction=LEFT, **kwargs):
+    def get_side_braces_and_labels(
+        self,
+        labels: list[str | VMobject | OpenGLVMobject],
+        direction: Vector3D = LEFT,
+        **kwargs: Any,
+    ) -> VGroup:
         assert hasattr(self, "horizontal_parts")
         parts = self.horizontal_parts
         return self.get_subdivision_braces_and_labels(
             parts, labels, direction, **kwargs
         )
 
-    def get_top_braces_and_labels(self, labels, **kwargs):
+    def get_top_braces_and_labels(
+        self, labels: list[str | VMobject | OpenGLVMobject], **kwargs: Any
+    ) -> VGroup:
         assert hasattr(self, "vertical_parts")
         parts = self.vertical_parts
         return self.get_subdivision_braces_and_labels(parts, labels, UP, **kwargs)
 
-    def get_bottom_braces_and_labels(self, labels, **kwargs):
+    def get_bottom_braces_and_labels(
+        self, labels: list[str | VMobject | OpenGLVMobject], **kwargs: Any
+    ) -> VGroup:
         assert hasattr(self, "vertical_parts")
         parts = self.vertical_parts
         return self.get_subdivision_braces_and_labels(parts, labels, DOWN, **kwargs)
 
-    def add_braces_and_labels(self):
+    def add_braces_and_labels(self) -> None:
         for attr in "horizontal_parts", "vertical_parts":
             if not hasattr(self, attr):
                 continue
@@ -176,11 +207,13 @@ class SampleSpace(Rectangle):
                 if hasattr(parts, subattr):
                     self.add(getattr(parts, subattr))
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int) -> SampleSpace:
         if hasattr(self, "horizontal_parts"):
-            return self.horizontal_parts[index]
+            val: SampleSpace = self.horizontal_parts[index]
+            return val
         elif hasattr(self, "vertical_parts"):
-            return self.vertical_parts[index]
+            val = self.vertical_parts[index]
+            return val
         return self.split()[index]
 
 
@@ -252,9 +285,8 @@ class BarChart(Axes):
         bar_width: float = 0.6,
         bar_fill_opacity: float = 0.7,
         bar_stroke_width: float = 3,
-        **kwargs,
+        **kwargs: Any,
     ):
-
         if isinstance(bar_colors, str):
             logger.warning(
                 "Passing a string to `bar_colors` has been deprecated since v0.15.2 and will be removed after v0.17.0, the parameter must be a list.  "
@@ -311,7 +343,7 @@ class BarChart(Axes):
 
         self.y_axis.add_numbers()
 
-    def _update_colors(self):
+    def _update_colors(self) -> None:
         """Initialize the colors of the bars of the chart.
 
         Sets the color of ``self.bars`` via ``self.bar_colors``.
@@ -319,17 +351,16 @@ class BarChart(Axes):
         Primarily used when the bars are initialized with ``self._add_bars``
         or updated via ``self.change_bar_values``.
         """
-
         self.bars.set_color_by_gradient(*self.bar_colors)
 
-    def _add_x_axis_labels(self):
+    def _add_x_axis_labels(self) -> None:
         """Essentially :meth`:~.NumberLine.add_labels`, but differs in that
         the direction of the label with respect to the x_axis changes to UP or DOWN
         depending on the value.
 
         UP for negative values and DOWN for positive values.
         """
-
+        assert isinstance(self.bar_names, list)
         val_range = np.arange(
             0.5, len(self.bar_names), 1
         )  # 0.5 shifted so that labels are centered, not on ticks
@@ -339,11 +370,8 @@ class BarChart(Axes):
         for i, (value, bar_name) in enumerate(zip(val_range, self.bar_names)):
             # to accommodate negative bars, the label may need to be
             # below or above the x_axis depending on the value of the bar
-            if self.values[i] < 0:
-                direction = UP
-            else:
-                direction = DOWN
-            bar_name_label = self.x_axis.label_constructor(bar_name)
+            direction = UP if self.values[i] < 0 else DOWN
+            bar_name_label: MathTex = self.x_axis.label_constructor(bar_name)
 
             bar_name_label.font_size = self.x_axis.font_size
             bar_name_label.next_to(
@@ -372,7 +400,6 @@ class BarChart(Axes):
         Rectangle
             A positioned rectangle representing a bar on the chart.
         """
-
         # bar measurements relative to the axis
 
         # distance from between the y-axis and the top of the bar
@@ -401,11 +428,11 @@ class BarChart(Axes):
 
     def get_bar_labels(
         self,
-        color: Color | None = None,
+        color: ParsableManimColor | None = None,
         font_size: float = 24,
         buff: float = MED_SMALL_BUFF,
-        label_constructor: type[VMobject] = Tex,
-    ):
+        label_constructor: type[MathTex] = Tex,
+    ) -> VGroup:
         """Annotates each bar with its corresponding value. Use ``self.bar_labels`` to access the
         labels after creation.
 
@@ -435,10 +462,9 @@ class BarChart(Axes):
 
                     self.add(chart, c_bar_lbls)
         """
-
         bar_labels = VGroup()
         for bar, value in zip(self.bars, self.values):
-            bar_lbl = label_constructor(str(value))
+            bar_lbl: MathTex = label_constructor(str(value))
 
             if color is None:
                 bar_lbl.set_color(bar.get_fill_color())
@@ -453,7 +479,9 @@ class BarChart(Axes):
 
         return bar_labels
 
-    def change_bar_values(self, values: Iterable[float], update_colors: bool = True):
+    def change_bar_values(
+        self, values: Iterable[float], update_colors: bool = True
+    ) -> None:
         """Updates the height of the bars of the chart.
 
         Parameters
@@ -483,7 +511,6 @@ class BarChart(Axes):
                     chart.change_bar_values(list(reversed(values)))
                     self.add(chart.get_bar_labels(font_size=24))
         """
-
         for i, (bar, value) in enumerate(zip(self.bars, values)):
             chart_val = self.values[i]
 
@@ -498,7 +525,6 @@ class BarChart(Axes):
             if chart_val != 0:
                 quotient = value / chart_val
                 if quotient < 0:
-
                     aligned_edge = UP if chart_val > 0 else DOWN
 
                     # if the bar is already positive, then we now want to move it
@@ -521,4 +547,4 @@ class BarChart(Axes):
         if update_colors:
             self._update_colors()
 
-        self.values[: len(values)] = values
+        self.values[: len(list(values))] = values

manim/mobject/graphing/scale.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, Any, Iterable
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any, overload
 
 import numpy as np
 
@@ -10,21 +11,29 @@ __all__ = ["LogBase", "LinearBase"]
 from manim.mobject.text.numbers import Integer
 
 if TYPE_CHECKING:
-    from manim.mobject.mobject import Mobject
+    from typing import Callable
+
+    from manim.mobject.types.vectorized_mobject import VMobject
 
 
 class _ScaleBase:
-    """Scale baseclass for graphing/functions."""
+    """Scale baseclass for graphing/functions.
+
+    Parameters
+    ----------
+    custom_labels
+        Whether to create custom labels when plotted on a :class:`~.NumberLine`.
+    """
 
     def __init__(self, custom_labels: bool = False):
-        """
-        Parameters
-        ----------
-        custom_labels
-            Whether to create custom labels when plotted on a :class:`~.NumberLine`.
-        """
         self.custom_labels = custom_labels
 
+    @overload
+    def function(self, value: float) -> float: ...
+
+    @overload
+    def function(self, value: np.ndarray) -> np.ndarray: ...
+
     def function(self, value: float) -> float:
         """The function that will be used to scale the values.
 
@@ -58,7 +67,8 @@ class _ScaleBase:
     def get_custom_labels(
         self,
         val_range: Iterable[float],
-    ) -> Iterable[Mobject]:
+        **kw_args: Any,
+    ) -> Iterable[VMobject]:
         """Custom instructions for generating labels along an axis.
 
         Parameters
@@ -89,7 +99,6 @@ class LinearBase(_ScaleBase):
         scale_factor
             The slope of the linear function, by default 1.0
         """
-
         super().__init__()
         self.scale_factor = scale_factor
 
@@ -124,7 +133,7 @@ class LogBase(_ScaleBase):
             The base of the log, by default 10.
         custom_labels
             For use with :class:`~.Axes`:
-            Whetherer or not to include ``LaTeX`` axis labels, by default True.
+            Whether or not to include ``LaTeX`` axis labels, by default True.
 
         Examples
         --------
@@ -139,13 +148,19 @@ class LogBase(_ScaleBase):
 
     def function(self, value: float) -> float:
         """Scales the value to fit it to a logarithmic scale.``self.function(5)==10**5``"""
-        return self.base**value
+        return_value: float = self.base**value
+        return return_value
 
     def inverse_function(self, value: float) -> float:
         """Inverse of ``function``. The value must be greater than 0"""
         if isinstance(value, np.ndarray):
             condition = value.any() <= 0
-            func = lambda value, base: np.log(value) / np.log(base)
+
+            func: Callable[[float, float], float]
+
+            def func(value: float, base: float) -> float:
+                return_value: float = np.log(value) / np.log(base)
+                return return_value
         else:
             condition = value <= 0
             func = math.log
@@ -161,8 +176,8 @@ class LogBase(_ScaleBase):
         self,
         val_range: Iterable[float],
         unit_decimal_places: int = 0,
-        **base_config: dict[str, Any],
-    ) -> list[Mobject]:
+        **base_config: Any,
+    ) -> list[Integer]:
         """Produces custom :class:`~.Integer` labels in the form of ``10^2``.
 
         Parameters
@@ -174,12 +189,11 @@ class LogBase(_ScaleBase):
         base_config
             Additional arguments to be passed to :class:`~.Integer`.
         """
-
         # uses `format` syntax to control the number of decimal places.
-        tex_labels = [
+        tex_labels: list[Integer] = [
             Integer(
                 self.base,
-                unit="^{%s}" % (f"{self.inverse_function(i):.{unit_decimal_places}f}"),
+                unit="^{%s}" % (f"{self.inverse_function(i):.{unit_decimal_places}f}"),  # noqa: UP031
                 **base_config,
             )
             for i in val_range