Allow non-numerical values to be added to :class:~.NumberLine/:class:~.Axes (#1780)

* add add_labels for number_line

* fix bug in add_numbers

* change to self

* drop add_coordinate_labels

* add somed docs

* improve docs

* forgot :class:

* fix typing

* remove redundancies

* typing error

* dict --> Dict

* introduce scaling in `add_labels` 

quite annoying to manually adjust size (especially without `font_size`)

* added tests

* don't use Text in testing

Author: hydrobeam

manim/mobject/coordinate_systems.py

@@ -12,7 +12,7 @@
 
 import fractions as fr
 import numbers
-from typing import Callable, Iterable, Optional, Sequence, Tuple, Union
+from typing import Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 import numpy as np
 from colour import Color
@@ -276,7 +276,11 @@ def get_axis_labels(
         return self.axis_labels
 
     def add_coordinates(
-        self, *axes_numbers: Optional[Iterable[float]], **kwargs
+        self,
+        *axes_numbers: Union[
+            Optional[Iterable[float]], Union[Dict[float, Union[str, float, "Mobject"]]]
+        ],
+        **kwargs,
     ) -> VGroup:
         """Adds labels to the axes.
 
@@ -294,6 +298,17 @@ def add_coordinates(
             ax.add_coordinates(x_labels, None, z_labels)  # default y labels, custom x & z labels
             ax.add_coordinates(x_labels)  # only x labels
 
+        ..code-block:: python
+
+            # specifically control the position and value of the labels using a dict
+            ax = Axes(x_range=[0, 7])
+            x_pos = [x for x in range(1, 8)]
+
+            # strings are automatically converted into a `Tex` mobject.
+            x_vals = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
+            x_dict = dict(zip(x_pos, x_vals))
+            ax.add_coordinates(x_dict)
+
         Returns
         -------
         VGroup
@@ -306,10 +321,13 @@ def add_coordinates(
             axes_numbers = [None for _ in range(self.dimension)]
 
         for axis, values in zip(self.axes, axes_numbers):
-            labels = axis.add_numbers(values, **kwargs)
+            if isinstance(values, dict):
+                labels = axis.add_labels(values, **kwargs)
+            else:
+                labels = axis.add_numbers(values, **kwargs)
             self.coordinate_labels.add(labels)
 
-        return self.coordinate_labels
+        return self
 
     def get_line_from_axis_to_point(
         self,

manim/mobject/number_line.py

@@ -3,9 +3,12 @@
 __all__ = ["NumberLine", "UnitInterval", "NumberLineOld"]
 
 import operator as op
+from typing import TYPE_CHECKING, Dict, Union
 
 import numpy as np
 
+from manim.mobject.svg.tex_mobject import MathTex, Tex
+
 from .. import config
 from ..constants import *
 from ..mobject.geometry import Line
@@ -18,6 +21,9 @@
 from ..utils.simple_functions import fdiv
 from ..utils.space_ops import normalize
 
+if TYPE_CHECKING:
+    from manim.mobject.mobject import Mobject
+
 
 class NumberLine(Line):
     """Creates a number line with tick marks. Number ranges that include both negative and
@@ -340,7 +346,51 @@ def add_numbers(self, x_values=None, excluding=None, **kwargs):
             numbers.add(self.get_number_mobject(x, **kwargs))
         self.add(numbers)
         self.numbers = numbers
-        return numbers
+        return self
+
+    def add_labels(
+        self,
+        dict_values: Dict[float, Union[str, float, "Mobject"]],
+        direction=None,
+        buff=None,
+    ):
+        """Adds specifically positioned labels to the :class:`~.NumberLine` using a ``dict``."""
+        if direction is None:
+            direction = self.label_direction
+        if buff is None:
+            buff = self.line_to_number_buff
+
+        labels = VGroup()
+        for x, label in dict_values.items():
+
+            label = self.create_label_tex(label)
+            label.scale(self.number_scale_value)
+            label.next_to(self.number_to_point(x), direction=direction, buff=buff)
+            labels.add(label)
+
+        self.labels = labels
+        self.add(labels)
+        return self
+
+    @staticmethod
+    def create_label_tex(label_tex) -> "Mobject":
+        """Checks if the label is a ``float``, ``int`` or a ``str`` and creates a :class:`~.MathTex`/:class:`~.Tex` label accordingly.
+
+        Parameters
+        ----------
+        label_tex : The label to be compared against the above types.
+
+        Returns
+        -------
+        :class:`~.Mobject`
+            The label.
+        """
+
+        if isinstance(label_tex, float) or isinstance(label_tex, int):
+            label_tex = MathTex(label_tex)
+        elif isinstance(label_tex, str):
+            label_tex = Tex(label_tex)
+        return label_tex
 
     def decimal_places_from_step(self):
         step_as_str = str(self.x_step)

tests/control_data/graphical_units_data/plot/custom_coordinates.npz

@@ -16,3 +16,13 @@ def test_axes(scene):
     )
     labels = graph.get_axis_labels()
     scene.add(graph, labels)
+
+
+@frames_comparison
+def test_custom_coordinates(scene):
+    ax = Axes(x_range=[0, 10])
+
+    ax.add_coordinates(
+        dict(zip([x for x in range(1, 10)], [Tex("str") for _ in range(1, 10)]))
+    )
+    scene.add(ax)

tests/test_graphical_units/test_axes.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from manim import NumberLine
+from manim.mobject.numbers import Integer
 
 
 def test_unit_vector():
@@ -50,3 +51,15 @@ def test_whole_numbers_step_size_default_to_0_decimal_places():
     assert actual_decimal_places == expected_decimal_places, (
         "Expected 1 decimal place but got " + actual_decimal_places
     )
+
+
+def test_add_labels():
+    expected_label_length = 6
+    num_line = NumberLine(x_range=[-4, 4])
+    num_line.add_labels(
+        dict(zip([x for x in range(-3, 3)], [Integer(m) for m in range(-1, 5)]))
+    )
+    actual_label_length = len(num_line.labels)
+    assert (
+        actual_label_length == expected_label_length
+    ), f"Expected a VGroup with {expected_label_length} integers but got {actual_label_length}."