Skip to content

Commit

Permalink
add some operators and utility funcs to BaseVector
Browse files Browse the repository at this point in the history
  • Loading branch information
duncathan committed Dec 13, 2023
1 parent c9c1c84 commit 1b404f5
Showing 1 changed file with 49 additions and 0 deletions.
49 changes: 49 additions & 0 deletions src/retro_data_structures/properties/base_vector.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -31,3 +33,50 @@ 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:
if center is None:
center = BaseVector()

pos = [self.x - center.x, self.y - center.y, self.z - center.z]

for i in range(3):
theta = rotation[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

0 comments on commit 1b404f5

Please sign in to comment.