A vector class with overloaded operators.

vector.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-

import sys
import math
import operator
import itertools

def pad(iterable, padding=0):
	for itm in itertools.chain(iterable, itertools.cycle([padding])):
		yield itm

class Vector(object):
	__slots__ = ['x', 'y', 'z', 'w']
	
	def __init__(self, *args):
		self.x = 0
		self.y = 0
		self.z = 0
		self.w = 0
		if len(args) == 1:
			args = args[0]
		for idx, arg in enumerate(args):
			self[idx] = arg
		
	def __getitem__(self, key):
		return self.__getattribute__(Vector.__slots__[key])
	def __setitem__(self, key, value):
		self.__setattr__(Vector.__slots__[key], value)
		return self
		
	def __getstate__(self):
		return [self.x, self.y, self.z, self.w]
	def __setstate__(self, state):
		self.x, self.y, self.z, self.w = state
		return self
	
	def __len__(self):
		return 4
		
	def __iter__(self):
		for i in xrange(4):
			yield self[i]
		raise StopIteration()
	
	def __repr__(self):
		return "Vector({}, {}, {}, {})".format(self.x, self.y, self.z, self.w)
	def __str__(self):
		return "Vector({}, {}, {}, {})".format(self.x, self.y, self.z, self.w)
	
	def __eq__(self, other):
		if hasattr(other, "__getitem__") and len(other) <= 4:
			return all(itertools.imap(operator.eq, self, pad(other, 0)))
		else:
			return False
	
	def __ne__(self, other):
		return not self.__eq__(other)
		
	def __op(self, other, f):
		"Any two-operator operation where the left operand is a Vector"
		if hasattr(other, "__getitem__"):
			return Vector(itertools.imap(f, self, pad(other)))
		else:
			return Vector(itertools.imap(f, self, itertools.cycle([other])))
 
	def __rop(self, other, f):
		"Any two-operator operation where the right operand is a Vector"
		if hasattr(other, "__getitem__"):
			return Vector(itertools.imap(f, pad(other), self))
		else:
			return Vector(itertools.imap(f, itertools.cycle([other]), self))
 
	def __iop(self, other, f):
		"inplace operator"
		if (hasattr(other, "__getitem__")):
			self.__setstate__(itertools.imap(f, self, pad(other)))
		else:
			self.__setstate__(itertools.imap(f, self, itertools.cycle([other])))
		return self
	
	def __add__(self, other):
		return self.__op(other, operator.add)
	def __radd__(self, other):
		return self.__rop(other, operator.add)
	def __iadd__(self, other):
		return self.__iop(other, operator.add)
	
	def __sub__(self, other):
		return self.__op(other, operator.sub)
	def __rsub__(self, other):
		return self.__rop(other, operator.sub)
	def __isub__(self, other):
		return self.__iop(other, operator.sub)
	
	def __mul__(self, other):
		return self.__op(other, operator.mul)
	def __rmul__(self, other):
		return self.__rop(other, operator.mul)
	def __imul__(self, other):
		return self.__iop(other, operator.mul)
	
	def __div__(self, other):
		return self.__op(other, operator.div)
	def __rdiv__(self, other):
		return self.__rop(other, operator.div)
	def __idiv__(self, other):
		return self.__iop(other, operator.div)
	
	def __truediv__(self, other):
		return self.__op(other, operator.truediv)
	def __rtruediv__(self, other):
		return self.__rop(other, operator.truediv)
	def __itruediv__(self, other):
		return self.__iop(other, operator.truediv)
	
	def __floordiv__(self, other):
		return self.__op(other, operator.floordiv)
	def __rfloordiv__(self, other):
		return self.__rop(other, operator.floordiv)
	def __ifloordiv__(self, other):
		return self.__iop(other, operator.floordiv)
	
	def __mod__(self, other):
		return self.__op(other, operator.mod)
	def __rmod__(self, other):
		return self.__rop(other, operator.mod)
		
	def __divmod__(self, other):
		return self.__op(other, operator.divmod)
	def __rdivmod__(self, other):
		return self.__rop(other, operator.divmod)
	
	# Exponentation
	def __pow__(self, other):
		return self.__op(other, operator.pow)
	def __rpow__(self, other):
		return self.__rop(other, operator.pow)
 
	# Bitwise operators
	def __lshift__(self, other):
		return self.__op(other, operator.lshift)
	def __rlshift__(self, other):
		return self.__rop(other, operator.lshift)
 
	def __rshift__(self, other):
		return self.__op(other, operator.rshift)
	def __rrshift__(self, other):
		return self.__rop(other, operator.rshift)
 
	def __and__(self, other):
		return self.__op(other, operator.and_)
	def __rand(self, other):
		return self.__rop(other, operator.and_)
 
	def __or__(self, other):
		return self.__op(other, operator.or_)
	def __ror__(self, other):
		return self.__rop(other, operator.or_)
 
	def __xor__(self, other):
		return self.__op(other, operator.xor)
	def __rxor__(self, other):
		return self.__rop(other, operator.xor)
	
 
	# Unary operations
	def __neg__(self):
		return Vector(*map(neg, self))
 
	def __pos__(self):
		return Vector(*map(pos, self))
 
	def __abs__(self):
		return Vector(*map(abs, self))
 
	def __invert__(self):
		return Vector(*map(operator.invert, self))
 
	def __int__(self):
		return Vector(*map(int, self))
 
	def __float__(self):
		return Vector(*map(float, self))
	
	def get_length_sqrd(self): 
		return sum(itertools.imap(operator.pow, self, itertools.cycle([2])))
 
	def get_length(self):
		return math.sqrt(sum(itertools.imap(operator.pow, self, itertools.cycle([2]))))
		
	def __setlength(self, value):
		length = self.get_length()
		self *= value/length
	length = property(get_length, __setlength, None, "gets or sets the magnitude of the vector")
	
	def normalized(self):
		length = self.length
		if length != 0:
			return self/length
		return Vector(self)
 
	def normalize_return_length(self):
		length = self.length
		if length != 0:
			self /= length
		return length
		
	def dot(self, other):
		return sum(itertools.imap(operator.mul, self, other))
		
	def cross(self, other):
		assert(len(other) >= 3)
		return Vector(
			self[1]*other[2] - self[2]*other[1],
			self[0]*other[2] - self[2]*other[0],
			self[1]*other[0] - self[0]*other[1]
		)
		
	def get_distance(self, other):
		return (self - other).length
		
	def get_dist_sqrd(self, other):
		return (self - other).get_length_sqrd()
	
class Spherical(object):
	__slots__ = ['r', 'theta', 'phi']
	def __init__(self, r = 1.0, theta = 0.0, phi = 0.0, vector = None):
		self.r     = r
		self.theta = theta
		self.phi   = phi
		if vector is not None:
			vec        = Vector(*vector)
			self.r     = vec.length
			self.theta = math.acos(vec.z / vec.length)
			self.phi   = math.atan2(vec.y, vec.x)
		
	def toCartesian(self):
		return Vector(
			r*math.sin(self.theta)*math.cos(self.phi),
			r*math.sin(self.theta)*math.sin(self.phi),
			r*math.cos(self.theta)
		)
		
class Plane(object):
	__slots__ = ['normal', 'point']
		
	def __init__(self, normal, point):
		self.normal = Vector(*normal)
		self.point  = Vector(*point)
		
	def __str__(self):
		return "Plane({},{})".format(self.normal, self.point)
		
	__repr__ = __str__
		
	def __contains__(self, point):
		return self.normal.dot(point - self.point) == 0
		
	def intersect(self, pos, vec):
		u   = Vector(*vec)
		w   = pos-self.point
		dot = self.normal.dot(vec)
		if abs(dot) > sys.float_info.epsilon:
			return pos + (u * -self.normal.dot(w) / dot)
		else:
			return None