add some operators and utility funcs to BaseVector

src/retro_data_structures/properties/base_vector.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import copy
 import dataclasses
+import math
 import typing
 
 from retro_data_structures.properties.base_property import BaseProperty
@@ -31,3 +33,59 @@ def to_json(self) -> json_util.JsonObject:
 
     def dependencies_for(self, asset_manager):
         yield from []
+
+    def __add__(self, other: BaseVector) -> typing_extensions.Self:
+        return self.__class__(self.x + other.x, self.y + other.y, self.z + other.z)
+
+    def __sub__(self, other: BaseVector) -> typing_extensions.Self:
+        return self.__class__(self.x - other.x, self.y - other.y, self.z - other.z)
+
+    def __mul__(self, other: int | float | BaseVector) -> typing_extensions.Self:
+        if isinstance(other, BaseVector):
+            return self.__class__(self.x * other.x, self.y * other.y, self.z * other.z)
+        if isinstance(other, int | float):
+            return self.__class__(self.x * other, self.y * other, self.z * other)
+        raise TypeError
+
+    def __truediv__(self, other: int | float | BaseVector) -> typing_extensions.Self:
+        if isinstance(other, BaseVector):
+            return self.__class__(self.x / other.x, self.y / other.y, self.z / other.z)
+        if isinstance(other, int | float):
+            return self.__class__(self.x / other, self.y / other, self.z / other)
+        raise TypeError
+
+    def __floordiv__(self, other: int | float | BaseVector) -> typing_extensions.Self:
+        if isinstance(other, BaseVector):
+            return self.__class__(self.x // other.x, self.y // other.y, self.z // other.z)
+        if isinstance(other, int | float):
+            return self.__class__(self.x // other, self.y // other, self.z // other)
+        raise TypeError
+
+    def rotate(self, rotation: BaseVector, center: BaseVector | None = None) -> typing_extensions.Self:
+        """
+        Rotates the vector on all three axes, around a center point.
+
+        :param rotation: The angle (in degrees) to rotate around each axis
+        :param center: The point around which to revolve
+        :returns: A new vector with the rotation applied
+        """
+
+        if center is None:
+            center = BaseVector()
+
+        pos = [self.x - center.x, self.y - center.y, self.z - center.z]
+        rot = [rotation.x, rotation.y, rotation.z]
+
+        for i in range(3):
+            theta = rot[i] * math.pi / 180.0
+            sin = math.sin(theta)
+            cos = math.cos(theta)
+
+            old_pos = copy.copy(pos)
+
+            comp1 = (i + 1) % 3
+            comp2 = (i + 2) % 3
+            pos[comp1] = old_pos[comp1] * cos - old_pos[comp2] * sin
+            pos[comp2] = old_pos[comp1] * sin + old_pos[comp2] * cos
+
+        return self.__class__(pos[0], pos[1], pos[2]) + center

tests/properties/test_base_vector.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import pytest
+
+from retro_data_structures.properties.base_vector import BaseVector
+
+
+def test_vector_operations():
+    a = BaseVector(2.0, 3.0, 4.0)
+    b = BaseVector(1.25, 1.5, 1.75)
+
+    assert a + b == BaseVector(3.25, 4.5, 5.75)
+    assert a - b == BaseVector(0.75, 1.5, 2.25)
+    assert a * b == BaseVector(2.5, 4.5, 7.0)
+    assert a / b == BaseVector(1.6, 2.0, 4.0 / 1.75)
+    assert a // b == BaseVector(1, 2, 2)
+
+
+@pytest.mark.parametrize(
+    ("center", "expected"),
+    (
+        (BaseVector(0.0, 0.0, 0.0), BaseVector(x=-3.5355339059327378, y=0.7071067811865481, z=-0.9999999999999998)),
+        (BaseVector(1.0, -3.0, 2.0), BaseVector(x=-3.2426406871192857, y=-5.82842712474619, z=2.0000000000000004)),
+        (BaseVector(2.4, 0.3, -24.03), BaseVector(x=-17.91517782348951, y=18.211014767455254, z=-22.63)),
+    ),
+)
+def test_rotation(center: BaseVector, expected: BaseVector):
+    pos = BaseVector(1.0, 2.0, 3.0)
+    rot = BaseVector(45.0, 90.0, 180.0)
+
+    result = pos.rotate(rot, center)
+    assert result == expected