fix(math): optimize __init__ in Vectors

K0lb3 · Feb 4, 2025 · 7886686 · 7886686
1 parent 5acb89a
commit 7886686
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/UnityPy/math/Vector3.py b/UnityPy/math/Vector3.py
@@ -15,21 +15,22 @@ class Vector3:
     Z: float = 0.0
 
     def __init__(self, *args):
+        from .Vector4 import Vector4
         if len(args) == 1:
             args = args[0]
 
         if isinstance(args, Sequence):
             if len(args) == 3:  # args=(x, y, z)
                 self.X, self.Y, self.Z = args
-                return
-            if len(args) == 0:  # args=()
+            elif len(args) == 0:  # args=()
                 self.X = self.Y = self.Z = 0.0
-                return
-        else:  # dirty patch for Vector4
+            else:
+                raise TypeError("Invalid argument length for Vector3")
+        elif isinstance(args, Vector4):
+            # dirty patch for Vector4
             self.X, self.Y, self.Z = args.X, args.Y, args.Z
-            return
-
-        raise TypeError("Invalid arguments for Vector3")
+        else:
+            raise TypeError("If only 1 argument passed, it must be a sequence or Vector4")
 
     def __getitem__(self, index):
         return (self.X, self.Y, self.Z)[index]

diff --git a/UnityPy/math/Vector4.py b/UnityPy/math/Vector4.py
@@ -22,15 +22,15 @@ def __init__(self, *args):
         if isinstance(args, Sequence):
             if len(args) == 4:  # args=(x, y, z, w)
                 self.X, self.Y, self.Z, self.W = args
-                return
-            if len(args) == 2:  # args=(Vector3, w)
+            elif len(args) == 2:  # args=(Vector3, w)
                 self.X, self.Y, self.Z = args[0]
                 self.W = args[1]
-            if len(args) == 0:  # args=()
+            elif len(args) == 0:  # args=()
                 self.X = self.Y = self.Z = self.W = 0.0
-                return
-
-        raise TypeError("Invalid arguments for Vector4")
+            else:
+                raise TypeError("Invalid argument length for Vector4")
+        else:
+            raise TypeError("If only 1 argument passed, it must be a sequence")
 
     def __getitem__(self, index):
         return (self.X, self.Y, self.Z, self.W)[index]