Merge pull request #56 from ppb/vector-like

Allow Vector-likes in all operations

ppb_vector/vector2.py

@@ -1,7 +1,38 @@
+import typing
+import collections
 from math import acos, cos, degrees, hypot, radians, sin
 from numbers import Real
 from collections.abc import Sequence
 
+__all__ = 'Vector2',
+
+
+VectorLike = typing.Union[
+    'Vector2',
+    typing.List[Real],  # TODO: Length 2
+    typing.Tuple[Real, Real],
+    typing.Dict[str, Real],  # TODO: Length 2, keys 'x', 'y'
+]
+
+
+def is_vector_like(value: typing.Any) -> bool:
+    return isinstance(value, (Vector2, list, tuple, dict))
+
+
+_fakevector = collections.namedtuple('_fakevector', ['x', 'y'])
+
+def _mkvector(value, *, castto=_fakevector):
+    if isinstance(value, Vector2):
+        return value
+    # FIXME: Allow all types of sequences
+    elif isinstance(value, (list, tuple)) and len(value) == 2:
+        return castto(value[0], value[1])
+    # FIXME: Allow all types of mappings
+    elif isinstance(value, dict) and 'x' in value and 'y' in value and len(value) == 2:
+        return castto(value['x'], value['y'])
+    else:
+        raise ValueError(f"Cannot use {value} as a vector-like")
+
 
 class Vector2(Sequence):
 
@@ -10,46 +41,55 @@ def __init__(self, x: Real, y: Real):
         self.y = y
         self.length = hypot(x, y)
 
-    def __len__(self):
+    @classmethod
+    def convert(cls, value: VectorLike) -> 'Vector2':
+        """
+        Constructs a vector from a vector-like.
+        """
+        return _mkvector(value, castto=type(cls))
+
+    def __len__(self) -> int:
         return 2
 
-    def __add__(self, other):
-        t = type(other)
-        if isinstance(other, Vector2):
-            return type(self)(self.x + other.x, self.y + other.y)
-        elif isinstance(other, (list, tuple)) and len(other) == 2:
-            return Vector2(self.x + other[0], self.y + other[1])
-        elif isinstance(other, (dict, set)) and 'x' in other and 'y' in other and len(other) == 2:
-            return Vector2(self.x + other['x'], self.y + other['y'])
-        else:
+    def __add__(self, other: VectorLike) -> 'Vector2':
+        try:
+            other = _mkvector(other)
+        except ValueError:
             return NotImplemented
+        rtype = type(other) if isinstance(other, Vector2) else type(self)
+        return rtype(self.x + other.x, self.y + other.y)
 
-    def __sub__(self, other):
-        if isinstance(other, Vector2):
-            return Vector2(self.x - other.x, self.y - other.y)
-        elif isinstance(other, (list, tuple)) and len(other) == 2:
-            return Vector2(self.x - other[0], self.y - other[1])
-        elif isinstance(other, dict) and 'x' in other and 'y' in other and len(other) == 2:
-            return Vector2(self.x - other['x'], self.y - other['y'])
-        else:
+    def __sub__(self, other: VectorLike) -> 'Vector2':
+        try:
+            other = _mkvector(other)
+        except ValueError:
             return NotImplemented
-
-    def __mul__(self, other):
-        if isinstance(other, Vector2):
+        rtype = type(other) if isinstance(other, Vector2) else type(self)
+        return rtype(self.x - other.x, self.y - other.y)
+
+    def __mul__(self, other: VectorLike) -> 'Vector2':
+        if is_vector_like(other):
+            try:
+                other = _mkvector(other)
+            except ValueError:
+                return NotImplemented
             return self.x * other.x + self.y * other.y
         elif isinstance(other, Real):
             return Vector2(self.x * other, self.y * other)
         else:
             return NotImplemented
 
-    def __rmul__(self, other):
-        if isinstance(other, Real):
-            return Vector2(self.x * other, self.y * other)
+    def __rmul__(self, other: VectorLike) -> 'Vector2':
+        return self.__mul__(other)
 
-    def __xor__(self, other):
+    def __xor__(self, other: VectorLike) -> Real:
+        """
+        Computes the magnitude of the cross product
+        """
+        other = _mkvector(other)
         return self.x * other.y - self.y * other.x
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: typing.Union[str, int]) -> Real:
         if hasattr(item, '__index__'):
             item = item.__index__()
         if isinstance(item, str):
@@ -69,29 +109,32 @@ def __getitem__(self, item):
         else:
             raise TypeError
 
-    def __repr__(self):
-        return "Vector2({}, {})".format(self.x, self.y)
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.x}, {self.y})"
 
-    def __eq__(self, other):
-        if isinstance(other, Vector2):
+    def __eq__(self, other: VectorLike) -> bool:
+        if is_vector_like(other):
+            other = _mkvector(other)
             return self.x == other.x and self.y == other.y
         else:
             return False
 
-    def __ne__(self, other):
-        if isinstance(other, Vector2):
+    def __ne__(self, other: VectorLike) -> bool:
+        if is_vector_like(other):
+            other = _mkvector(other)
             return self.x != other.x or self.y != other.y
         else:
             return True
 
-    def __iter__(self):
+    def __iter__(self) -> typing.Iterator[Real]:
         yield self.x
         yield self.y
 
-    def __neg__(self):
+    def __neg__(self) -> 'Vector2':
         return self * -1
 
-    def angle(self, other):
+    def angle(self, other: VectorLike) -> Real:
+        other = _mkvector(other, castto=Vector2)
         return degrees(acos(self.normalize() * other.normalize()))
 
     def isclose(self, other: 'Vector2', *, rel_tol: float=1e-06, abs_tol: float=1e-3):
@@ -116,36 +159,37 @@ def isclose(self, other: 'Vector2', *, rel_tol: float=1e-06, abs_tol: float=1e-3
             diff < abs_tol
         )
 
-    def rotate(self, degrees):
+    def rotate(self, degrees: Real) -> 'Vector2':
         r = radians(degrees)
         r_cos = cos(r)
         r_sin = sin(r)
         x = round(self.x * r_cos - self.y * r_sin, 5)
         y = round(self.x * r_sin + self.y * r_cos, 5)
         return Vector2(x, y)
 
-    def normalize(self):
+    def normalize(self) -> 'Vector2':
         return self.scale(1)
 
-    def truncate(self, max_length):
+    def truncate(self, max_length: Real) -> 'Vector2':
         if self.length > max_length:
             return self.scale(max_length)
         return self
 
-    def scale(self, length):
+    def scale(self, length: Real) -> 'Vector2':
         try:
             scale = length / self.length
         except ZeroDivisionError:
             scale = 1
         return self * scale
 
-    def reflect(self, surface_normal: 'Vector2') -> 'Vector2':
+    def reflect(self, surface_normal: VectorLike) -> 'Vector2':
         """
         Calculate the reflection of the vector against a given surface normal
         """
+        surface_normal = _mkvector(surface_normal, castto=Vector2)
         if not (0.99999 < surface_normal.length < 1.00001):
             raise ValueError("Reflection requires a normalized vector.")
         vec_new = self
-        if self * surface_normal>0:
+        if self * surface_normal > 0:
             vec_new = self.rotate(180)
         return vec_new - (2 * (vec_new * surface_normal) * surface_normal)