Fixes for ComplexPlane

Author: JasonGrace2282

manim/mobject/graphing/coordinate_systems.py

@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
 
 import numpy as np
-from typing_extensions import Self
+from typing_extensions import Self, TypedDict
 
 from manim import config
 from manim.constants import *
@@ -55,14 +55,25 @@
 from manim.utils.space_ops import angle_of_vector
 
 if TYPE_CHECKING:
-    import numpy.typing as npt
-
     from manim.mobject.mobject import Mobject
     from manim.typing import ManimFloat, Point2D, Point3D, Vector3D
 
     LineType = TypeVar("LineType", bound=Line)
 
 
+class _MatmulConfig(TypedDict):
+    """A dictionary for configuring the __matmul__/__rmatmul__ operation.
+
+    Parameters
+    ----------
+        method: The method to call
+        unpack: whether to unpack the parameter given to __matmul__/__rmatmul__
+    """
+
+    method: str
+    unpack: bool
+
+
 class CoordinateSystem:
     r"""Abstract base class for Axes and NumberPlane.
 
@@ -1793,20 +1804,29 @@ def construct(self):
 
         return T_label_group
 
-    _matmul_method = "coords_to_point"
-    _rmatmul_method = "point_to_coords"
+    _matmul_config: _MatmulConfig = {
+        "method": "coords_to_point",
+        "unpack": True,
+    }
+    _rmatmul_config: _MatmulConfig = {"method": "point_to_coords", "unpack": False}
 
-    def __matmul__(self, coord: Sequence[float] | Mobject | npt.NDArray[np.float64]):
+    def __matmul__(self, coord):
         if isinstance(coord, Mobject):
             coord = coord.get_center()
-        method = getattr(self, self._matmul_method)
+        method = getattr(self, self._matmul_config["method"])
         assert callable(method)
-        return method(*coord)
+        return (
+            method(*coord) if self._matmul_config.get("unpack", True) else method(coord)
+        )
 
-    def __rmatmul__(self, point: Point3D):
-        method = getattr(self, self._rmatmul_method)
+    def __rmatmul__(self, point):
+        method = getattr(self, self._rmatmul_config["method"])
         assert callable(method)
-        return method(point)
+        return (
+            method(*point)
+            if self._rmatmul_config.get("unpack", False)
+            else method(point)
+        )
 
 
 class Axes(VGroup, CoordinateSystem, metaclass=ConvertToOpenGL):
@@ -2954,8 +2974,11 @@ def construct(self):
                 self.add(polarplane_pi)
     """
 
-    _matmul_method = "polar_to_point"
-    _rmatmul_method = "point_to_polar"
+    _matmul_config = {
+        "method": "polar_to_point",
+        "unpack": True,
+    }
+    _rmatmul_config = {"method": "point_to_polar", "unpack": False}
 
     def __init__(
         self,
@@ -3327,6 +3350,10 @@ def construct(self):
 
     """
 
+    _matmul_config = {"method": "number_to_point", "unpack": False}
+
+    _rmatmul_config = {"method": "point_to_number", "unpack": False}
+
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(
             **kwargs,

tests/module/mobject/graphing/test_coordinate_system.py

@@ -202,11 +202,13 @@ def test_matmul_operations():
     mob = Dot().move_to((1, 2, 0))
     assert (ax @ mob == ax.coords_to_point(1, 2)).all()
 
-    # other coordinate systems like PolarPlane should override __matmul__ indirectly
+    # other coordinate systems like PolarPlane and ComplexPlane should override __matmul__ indirectly
     polar = PolarPlane()
-    # radius, azimuthal angle
     assert (polar @ (1, 2) == polar.polar_to_point(1, 2)).all()
 
+    complx = ComplexPlane()
+    assert (complx @ (1 + 2j) == complx.number_to_point(1 + 2j)).all()
+
     # Numberline doesn't inherit from CoordinateSystem, but it should still work
     n = NumberLine()
     assert (n @ 3 == n.number_to_point(3)).all()