Refactor the math module

UnityPy/math/Color.py

@@ -1,24 +1,22 @@
+from dataclasses import dataclass
 from .Vector4 import Vector4
 
 
+@dataclass
 class Color:
     R: float
     G: float
     B: float
     A: float
 
     def __init__(self, r: float = 0.0, g: float = 0.0, b: float = 0.0, a: float = 0.0):
+        if not all(isinstance(v, (int, float)) for v in (r, g, b, a)):
+            raise TypeError("All components must be numeric.")
         self.R = r
         self.G = g
         self.B = b
         self.A = a
 
-    def __eq__(self, other):
-        if isinstance(other, Color):
-            return self.__dict__ == other.__dict__
-        else:
-            return False
-
     def __add__(self, other):
         return Color(
             self.R + other.R, self.G + other.G, self.B + other.B, self.A + other.A
@@ -37,7 +35,7 @@ def __mul__(self, other):
         else:
             return Color(self.R * other, self.G * other, self.B * other, self.A * other)
 
-    def __div__(self, other):
+    def __truediv__(self, other):
         if isinstance(other, Color):
             return Color(
                 self.R / other.R, self.G / other.G, self.B / other.B, self.A / other.A
@@ -46,10 +44,13 @@ def __div__(self, other):
             return Color(self.R / other, self.G / other, self.B / other, self.A / other)
 
     def __eq__(self, other):
-        return self.__dict__ == other.__dict__
+        if isinstance(other, Color):
+            return self.__dict__ == other.__dict__
+        else:
+            return False
 
     def __ne__(self, other):
-        return self.__dict__ != other.__dict__
+        return not (self == other)
 
     def Vector4(self):
         return Vector4(self.R, self.G, self.B, self.A)

UnityPy/math/Matrix4x4.py

@@ -1,31 +1,43 @@
+from dataclasses import dataclass
+from typing import MutableSequence, Sequence, Union
 from .Vector3 import Vector3
 
 
+@dataclass
 class Matrix4x4:
-    M: list
+    M: MutableSequence[float]
 
-    def __init__(self, values):
-        if len(values) != 16:
-            raise ValueError(
-                "There must be sixteen and only sixteen input values for Matrix."
-            )
-        self.M = values
+    def __init__(self, values: Sequence[Union[int, float]]):
+        if not isinstance(values, Sequence) or len(values) != 16:
+            raise ValueError("Values must be a sequence with 16 elements.")
+        if not all(isinstance(v, (int, float)) for v in values):
+            raise TypeError("All values must be numeric.")
+        self.M = [float(v) for v in values]
 
     def __getitem__(self, index):
         if isinstance(index, tuple):
-            index = index[0] + index[1] * 4
+            row, col = index
+            if not (0 <= row < 4 and 0 <= col < 4):
+                raise IndexError("Row and column indices must in range [0, 3].")
+            index = row + col * 4
+        if not (0 <= index < 16):
+            raise IndexError("Index out of range for Matrix4x4.")
         return self.M[index]
 
     def __setitem__(self, index, value):
         if isinstance(index, tuple):
-            # row, column
-            index = index[0] + index[1] * 4
+            row, col = index
+            if not (0 <= row < 4 and 0 <= col < 4):
+                raise IndexError("Row and column indices must in range [0, 3].")
+            index = row + col * 4
+        if not (0 <= index < 16):
+            raise IndexError("Index out of range for Matrix4x4.")
         self.M[index] = value
 
     def __eq__(self, other):
         if not isinstance(other, Matrix4x4):
             return False
-        print()
+        return all(abs(a - b) < 1e-6 for a, b in zip(self.M, other.M))
 
     def __mul__(lhs, rhs):
         res = Matrix4x4([0] * 16)
@@ -134,7 +146,10 @@ def __mul__(lhs, rhs):
     @staticmethod
     def Scale(vector: Vector3):
         return Matrix4x4(
-            [vector.X, 0, 0, 0, 0, vector.Y, 0, 0, 0, 0, vector.Z, 0, 0, 0, 0, 1]
+            [vector.X, 0, 0, 0,
+             0, vector.Y, 0, 0,
+             0, 0, vector.Z, 0,
+             0, 0, 0, 1]
         )
 
     @property

UnityPy/math/Quaternion.py

@@ -1,50 +1,32 @@
+from dataclasses import dataclass
+
+
+@dataclass
 class Quaternion:
     X: float
     Y: float
     Z: float
     W: float
-
-    def __init__(self, x: float = 0.0, y: float = 0.0, z: float = 0.0, w: float = 0.0):
-        self._data = [0.0] * 4
-        self.X = x
-        self.Y = y
-        self.Z = z
-        self.W = w
-
-    @property
-    def X(self) -> float:
-        return self._data[0]
-
-    @X.setter
-    def X(self, value: float):
-        self._data[0] = value
-
-    @property
-    def Y(self) -> float:
-        return self._data[1]
-
-    @Y.setter
-    def Y(self, value: float):
-        self._data[1] = value
-
-    @property
-    def Z(self) -> float:
-        return self._data[2]
-
-    @Z.setter
-    def Z(self, value: float):
-        self._data[2] = value
-
-    @property
-    def W(self) -> float:
-        return self._data[3]
-
-    @W.setter
-    def W(self, value: float):
-        self._data[3] = value
-
-    def __getitem__(self, value):
-        return self._data[value]
-
+
+    def __init__(self, x: float = 0.0, y: float = 0.0, z: float = 0.0, w: float = 1.0):
+        if not all(isinstance(v, (int, float)) for v in (x, y, z, w)):
+            raise TypeError("All components must be numeric.")
+        self.X = float(x)
+        self.Y = float(y)
+        self.Z = float(z)
+        self.W = float(w)
+
+    def __getitem__(self, index):
+        return (self.X, self.Y, self.Z, self.W)[index]
+
     def __setitem__(self, index, value):
-        self._data[index] = value
+        if index == 0:
+            self.X = value
+        elif index == 1:
+            self.Y = value
+        elif index == 2:
+            self.Z = value
+        elif index == 3:
+            self.W = value
+        else:
+            raise IndexError("Index out of range")

UnityPy/math/Rectangle.py

@@ -1,3 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
 class Rectangle:
     height: int
     width: int