vector.py

from math import sqrt
from operator import (add, sub, mul, div, floordiv, mod, pow as opow, truediv,
                      lt, le, eq, ne, gt, ge)

class VectorError(Exception): pass # base exception
class InvalidDimension(VectorError): pass
InvalidDimensionError = InvalidDimension

ELEMENTS = dict(zip("xyzw", xrange(4)))

class Vector(list):
    """
    Vector class meant to mimic classical mathematical notation

    The difference is that column vectors are defined instead as just lists or tuples
    """

    def __init__(self, *args):
        if len(args) == 1:
            args = args[0]
        list.__init__(self, args)

        if len(self) < 2: # no < 2D vectors
            self.extend([0]*(2-len(self)))
        self._assert_contents() # check for a valid vector

    def __repr__(self):
        return repr(tuple(map(float, self)))

    def _get_mag(self):
        return sqrt(sum(map(lambda x: x**2, self)))

    def _set_mag(self, value):
        self.normalize()
        self *= value
        return value

    length = magnitude = property(_get_mag, _set_mag, doc="The length of the vector")

    def _assert_contents(self):
        for i in xrange(len(self)):
            x = self[i]
            if not isinstance(x, (int, float, long, complex)):
                raise ValueError('Invalid Vector Value: (%r)'%x)
            elif isinstance(x, complex):
                pass
            else:
                self[i] = float(x)


    ### ==== LIST OPERATORS ==== ###

    def __setitem__(self, key, value):
        if isinstance(key, slice):
            stop = key.stop
        else:
            stop = key

        while len(self) < (stop+1):
            self.append(0)

        list.__setitem__(self, key, float(value))
        #self._assert_contents()

    ### ==== SWIZZLING ==== ###

    def __getattr__(self, name):
        if not len([x for x in name if x in ELEMENTS]) == len(name):
            raise AttributeError("Vector object has no attribute '%s'"%name)
        v = []
        for x in name:
            if len(self) <= ELEMENTS[x]:
                raise AttributeError("Vector object has no element '%s'"%x)
            v.append(self[ELEMENTS[x]])
        return Vector(v) if len(v) > 1 else v[0]

    def __setattr__(self, name, value):
        made_up_of_elements = len([x for x in name if x in ELEMENTS])==len(name)
        if not made_up_of_elements and name not in self.__dict__:
            raise AttributeError("Vector object has no attribute '%s'"%name)
        elif len(name) > 1 and (not hasattr(value, '__iter__') or len(name) != len(value)):
            raise TypeError("Vector size does not match swizzle")
        elif name in self.__dict__:
            self.__dict__[name] = value
        if not hasattr(value, "__iter__"):
            value = [value]
        for i in range(len(name)):
            vname, val = name[i], value[i]
            idx = ELEMENTS[vname]
            self[idx] = float(val)

    ### ==== OPERATORS ==== ###

    def _test_length(self, other):
        if not len(self) == len(other):
            raise InvalidDimensionError

    def _apply_op(self, other, op):
        vec = Vector()
        self._test_length(other)
        for i in xrange(len(self)):
            vec[i] = op(self[i], other[i])
        return vec

    def __add__(self, other):
        return self._apply_op(other, add)

    def __sub__(self, other):
        return self._apply_op(other, sub)

    def __mul__(self, other):
        if not isinstance(other, (int, float)):
            return self.dot(other)
        return Vector(map(lambda x: mul(x, other), self))

    def __div__(self, other):
        return Vector(map(lambda x: div(x, other), self))

    def __truediv__(self, other):
        return Vector(map(lambda x: truediv(x, other), self))

    def __floordiv__(self, other):
        return Vector(map(lambda x: floordiv(x, other), self))

    def __mod__(self, other):
        return Vector(map(lambda x: mod(x, other), self))

    def __pow__(self, other, modulo=None):
        v = Vector(map(lambda x: opow(x, other), self))
        if modulo is not None:
            v % modulo
        return v

    ## == Right Operators == ##

    def __rmul__(self, other):
        if not isinstance(other, (int, float)):
            return Vector(other).dot(self)
        return Vector(map(lambda x: mul(other, x), self))

    def __rdiv__(self, other):
        return Vector(map(lambda x: div(other, x), self))

    def __rtruediv__(self, other):
        return Vector(map(lambda x: truediv(other, x), self))

    def __rfloordiv__(self, other):
        return Vector(map(lambda x: floordiv(other, x), self))

    ## == Other Operators == ##

    def __iadd__(self, other):
        self = add(self, other)
        return self

    def __isub__(self, other):
        self = sub(self, other)
        return self

    def __imul__(self, other):
        self = mul(self, other)
        return self

    def __idiv__(self, other):
        self = div(self, other)
        return self

    def __itruediv__(self, other):
        self = truediv(self, other)
        return self

    def __ifloordiv__(self, other):
        self = floordiv(self, other)
        return self

    def __imod__(self, other):
        self = mod(self, other)
        return self

    def __ipow__(self, other, modulo=None):
        self = pow(self, other, modulo)
        return self

    def __neg__(self):
        return self * -1

    def __pos__(self):
        return self * +1

    def __abs__(self):
        return Vector(map(lambda x: abs(x), self))

    def __invert__(self):
        return 1.0 / self

    ## == Comparison Operators == ##

    def __lt__(self, other):
        return lt(self.magnitude, other.magnitude)

    def __le__(self, other):
        return le(self.magnitude, other.magnitude)

    def __eq__(self, other):
        return all(self._apply_op(other, eq))

    def __ne__(self, other):
        return all(self._apply_op(other, ne))

    def __gt__(self, other):
        return gt(self.magnitude, other.magnitude)

    def __ge__(self, other):
        return ge(self.magnitude, other.magnitude)

    ## == Vector Operations == ##

    def apply(self, func):
        return Vector(map(func, self))

    def copy(self):
        return Vector(iter(self))

    def zero(self):
        for i in xrange(len(self)):
            self[i] = 0.0
        return self

    def normalize(self):
        mag = self.magnitude
        for i in xrange(len(self)):
            self[i] /= mag
        return self

    def negate(self):
        self *= -1
        return self

    def reflect(self, mirror):
        self = mirror - (2 * self * self.dot(mirror))
        return self

    def _make_length(self, n):
        if len(self) < n:
            self.extend([0.0]*(n-len(self)))
            return True
        elif len(self) > n:
            del self[n:]
            return False

    def resize2D(self):
        self._make_length(2)
        return self

    def resize3D(self):
        self._make_length(3)
        return self

    def resize4D(self):
        if self._make_length(4):
            self[3] = 1.0
        return self

    def cross(self, other):
        pass

    def dot(self, other):
        self._test_length(other)
        return sum([self[i] * other[i] for i in xrange(len(self))], 0.0)

if __name__ == "__main__":
    v = Vector(1, 2, 3)
    u = Vector(5, 10, 15)