Switch arcade.math.lerp to use a Protocol bound to a TypeVar (#2310)

* Add HasAddSubMul protocol to cover lerping

* Add support for arbitary lerpable arguments to arcade.math.lerp

* Explain why the protocol and TypeVars are so specific

* Convert lerp_2d and lerp_3d to return Vec types

* Use appropriate Point2 and Point3 annotations instead of Point or local annotations

* Fix pyright by cleaning A* pathing annotations and imports

Author: pushfoo

arcade/math.py

@@ -2,11 +2,11 @@
 
 import math
 import random
-from typing import Sequence, Union
+from typing import TypeVar
 
-from pyglet.math import Vec2
+from pyglet.math import Vec2, Vec3
 
-from arcade.types import AsFloat, Point, Point2
+from arcade.types import HasAddSubMul, Point, Point2
 from arcade.types.rect import Rect
 from arcade.types.vector_like import Point3
 
@@ -46,43 +46,64 @@ def clamp(a, low: float, high: float) -> float:
     return high if a > high else max(a, low)
 
 
-V_2D = Union[tuple[AsFloat, AsFloat], Sequence[AsFloat]]
-V_3D = Union[tuple[AsFloat, AsFloat, AsFloat], Sequence[AsFloat]]
+# This TypeVar helps match v1 and v2 as the same type below in lerp's
+# signature. If we used HasAddSubMul, they could be different.
+L = TypeVar("L", bound=HasAddSubMul)
 
 
-def lerp(v1: AsFloat, v2: AsFloat, u: float) -> float:
-    """linearly interpolate between two values
+def lerp(v1: L, v2: L, u: float) -> L:
+    """Linearly interpolate two values which support arithmetic operators.
+
+    Both ``v1`` and ``v2`` must be of compatible types and support
+    the following operators:
+
+    * ``+`` (:py:meth:`~object.__add__`)
+    * ``-`` (:py:meth:`~object.__sub__`)
+    * ``*`` (:py:meth:`~object.__mul__`)
+
+    This means that in certain cases, you may want to use another
+    function:
+
+    * For angles, use :py:func:`lerp_angle`.
+    * To convert points as arbitary sequences, use:
+
+      * :py:func:`lerp_2d`
+      * :py:func:`lerp_3d`
 
     Args:
-        v1 (float): The first value
-        v2 (float): The second value
-        u (float): The interpolation value `(0.0 to 1.0)`
+        v1 (HasAddSubMul): The first value
+        v2 (HasAddSubMul): The second value
+        u: The interpolation value `(0.0 to 1.0)`
     """
     return v1 + ((v2 - v1) * u)
 
 
-def lerp_2d(v1: V_2D, v2: V_2D, u: float) -> tuple[float, float]:
-    """
-    Linearly interpolate between two 2D points.
+def lerp_2d(v1: Point2, v2: Point2, u: float) -> Vec2:
+    """Linearly interpolate between two 2D points passed as sequences.
+
+    .. tip:: This function returns a :py:class:`Vec2` you can use
+             with :py:func`lerp` .
 
     Args:
-        v1 (tuple[float, float]): The first point
-        v2 (tuple[float, float]): The second point
+        v1: The first point as a sequence of 2 values.
+        v2: The second point as a sequence of 2 values.
         u (float): The interpolation value `(0.0 to 1.0)`
     """
-    return (lerp(v1[0], v2[0], u), lerp(v1[1], v2[1], u))
+    return Vec2(lerp(v1[0], v2[0], u), lerp(v1[1], v2[1], u))
 
 
-def lerp_3d(v1: V_3D, v2: V_3D, u: float) -> tuple[float, float, float]:
-    """
-    Linearly interpolate between two 3D points.
+def lerp_3d(v1: Point3, v2: Point3, u: float) -> Vec3:
+    """Linearly interpolate between two 3D points passed as sequences.
+
+    .. tip:: This function returns a :py:class:`Vec2` you can use
+             with :py:func`lerp`.
 
     Args:
-        v1 (tuple[float, float, float]): The first point
-        v2 (tuple[float, float, float]): The second point
+        v1: The first point as a sequence of 3 values.
+        v2: The second point as a sequence of 3 values.
         u (float): The interpolation value `(0.0 to 1.0)`
     """
-    return (lerp(v1[0], v2[0], u), lerp(v1[1], v2[1], u), lerp(v1[2], v2[2], u))
+    return Vec3(lerp(v1[0], v2[0], u), lerp(v1[1], v2[1], u), lerp(v1[2], v2[2], u))
 
 
 def lerp_angle(start_angle: float, end_angle: float, u: float) -> float:

arcade/paths.py

@@ -5,11 +5,10 @@
 from __future__ import annotations
 
 import math
-from typing import cast
 
 from arcade import Sprite, SpriteList, check_for_collision_with_list, get_sprites_at_point
 from arcade.math import get_distance, lerp_2d
-from arcade.types import Point, Point2
+from arcade.types import Point2
 
 __all__ = ["AStarBarrierList", "astar_calculate_path", "has_line_of_sight"]
 
@@ -33,13 +32,13 @@ def _spot_is_blocked(position: Point2, moving_sprite: Sprite, blocking_sprites:
     return len(hit_list) > 0
 
 
-def _heuristic(start: Point, goal: Point) -> float:
+def _heuristic(start: Point2, goal: Point2) -> float:
     """
     Returns a heuristic value for the passed points.
 
     Args:
-        start (Point): The 1st point to compare
-        goal (Point): The 2nd point to compare
+        start (Point2): The 1st point to compare
+        goal (Point2): The 2nd point to compare
 
     Returns:
         float: The heuristic of the 2 points
@@ -102,7 +101,7 @@ def __init__(
         else:
             self.movement_directions = (1, 0), (-1, 0), (0, 1), (0, -1)  # type: ignore
 
-    def get_vertex_neighbours(self, pos: Point) -> list[tuple[float, float]]:
+    def get_vertex_neighbours(self, pos: Point2) -> list[tuple[float, float]]:
         """
         Return neighbors for this point according to ``self.movement_directions``
 
@@ -123,7 +122,7 @@ def get_vertex_neighbours(self, pos: Point) -> list[tuple[float, float]]:
             n.append((x2, y2))
         return n
 
-    def move_cost(self, a: Point, b: Point) -> float:
+    def move_cost(self, a: Point2, b: Point2) -> float:
         """
         Returns a float of the cost to move
 
@@ -224,12 +223,12 @@ def _AStarSearch(start: Point2, end: Point2, graph: _AStarGraph) -> list[Point2]
     return None
 
 
-def _collapse(pos: Point, grid_size: float):
+def _collapse(pos: Point2, grid_size: float) -> tuple[int, int]:
     """Makes Point pos smaller by grid_size"""
     return int(pos[0] // grid_size), int(pos[1] // grid_size)
 
 
-def _expand(pos: Point, grid_size: float):
+def _expand(pos: Point2, grid_size: float) -> tuple[int, int]:
     """Makes Point pos larger by grid_size"""
     return int(pos[0] * grid_size), int(pos[1] * grid_size)
 
@@ -329,11 +328,11 @@ def recalculate(self):
 
 
 def astar_calculate_path(
-    start_point: Point,
-    end_point: Point,
+    start_point: Point2,
+    end_point: Point2,
     astar_barrier_list: AStarBarrierList,
     diagonal_movement: bool = True,
-) -> list[Point] | None:
+) -> list[Point2] | None:
     """
     Calculates the path using AStarSearch Algorithm and returns the path
 
@@ -371,13 +370,13 @@ def astar_calculate_path(
 
     # Currently 'result' is in grid locations. We need to convert them to pixel
     # locations.
-    revised_result = [_expand(p, grid_size) for p in result]
-    return cast(list[Point], revised_result)
+    revised_result: list[Point2] = [_expand(p, grid_size) for p in result]
+    return revised_result
 
 
 def has_line_of_sight(
-    observer: Point,
-    target: Point,
+    observer: Point2,
+    target: Point2,
     walls: SpriteList,
     max_distance: float = float("inf"),
     check_resolution: int = 2,
@@ -429,7 +428,7 @@ def has_line_of_sight(
 
 
 # NOTE: Rewrite this
-# def dda_step(start: Point, end: Point):
+# def dda_step(start: Point2, end: Point2):
 #     """
 #     Bresenham's line algorithm
 

arcade/types/__init__.py

@@ -26,7 +26,7 @@
 # flake8: noqa: E402
 import sys
 from pathlib import Path
-from typing import NamedTuple, Union, TYPE_CHECKING, TypeVar, Iterable
+from typing import NamedTuple, Union, TYPE_CHECKING, TypeVar, Iterable, Protocol
 
 from pytiled_parser import Properties
 
@@ -124,6 +124,7 @@
     "Box",
     "LRBTNF",
     "XYZWHD",
+    "HasAddSubMul",
     "RGB",
     "RGBA",
     "RGBOrA",
@@ -206,6 +207,23 @@ def annotated2(argument: OneOrIterableOf[MyType] | None = tuple()):
 # --- End potentially obsolete annotations ---
 
 
+# These are for the argument type + return type. They're separate TypeVars
+# to handle cases which take tuple but return Vec2 (e.g. pyglet.math.Vec2).
+_T_contra = TypeVar("_T_contra", contravariant=True)  # Same or more general than T
+_T_co = TypeVar("_T_co", covariant=True)  # Same or more specific than T
+
+
+class HasAddSubMul(Protocol[_T_contra, _T_co]):
+    """Matches types which work with :py:func:`arcade.math.lerp`."""
+
+    # The / matches float and similar operations to keep pyright
+    # happy since built-in arithmetic makes them positional only.
+    # See https://peps.python.org/pep-0570/
+    def __add__(self, value: _T_contra, /) -> _T_co: ...
+    def __sub__(self, value: _T_contra, /) -> _T_co: ...
+    def __mul__(self, value: _T_contra, /) -> _T_co: ...
+
+
 # Path handling
 PathLike = Union[str, Path, bytes]
 _POr = TypeVar("_POr")  # Allows PathOr[TypeNameHere] syntax