Example Coordinate class for Pandas Extension Types article

vector.py

from functools import total_ordering
from math import sqrt

@total_ordering
class Vector(object):
    """
    Simple class to represent a 2D vector with X and Y components.
    Could extend with more useful methods etc
    """
    def __init__(self, x, y):
        self.x = float(x)
        self.y = float(y)

    def __getitem__(self, index):
        """
        Allows object to act like (x, y) pair with indexing
        """
        if index == 0:
            return self.x
        elif index == 1:
            return self.y
        else:
            raise KeyError('Invalid vector index: {}'.format(index))

    def as_tuple(self):
        """
        Return as (x, y) vector pair
        """
        return (self.x, self.y)

    def magnitude(self):
        """
        Get magnitude of vector
        """
        return sqrt(self.x**2 + self.y**2)
        
    def dot(self, other):
        """
        Calculate dot product with another Vector
        """
        other_vector = create_vector(other)
        return self.x*other_vector.x + self.y*other_vector.y

    def __len__(self):
        return 2

    def __repr__(self):
        return 'Vector({}, {})'.format(self.x, self.y)

    # Operator support
    def __add__(self, other):
        """
        Vector Addition
        """
        other_vector = create_vector(other)
        return Vector(self.x + other_vector.x, self.y + other_vector.y)
    
    def __sub__(self, other):
        """
        Vector Subtraction
        """
        other_vector = create_vector(other)
        return Vector(self.x - other_vector.x, self.y - other_vector.y)

    def __mul__(self, other):
        if isinstance(other, (int, float)):
            return Vector(self.x * other, self.y * other)
        else:
            raise TypeError('Cannot multiply vector by {}'.format(type(other)))

    def __neg__(self):
        return Vector(-self.x, -self.y)

    def __eq__(self, other):
        other_coord = create_vector(other)
        return self.x == other_coord.x and self.y == other_coord.y

    def __lt__(self, other):
        other_coord = create_vector(other)
        return self.x < other_coord.x and self.y < other_coord.y 


def create_vector(val):
    """
    Factory function for constructing a Vector from various
    types of inputs
    """
    if isinstance(val, Vector):
        return val

    if isinstance(val, (list, tuple)) and len(val) == 2:
        # Construct from list-like of X,Y value pair
        return Vector(val[0], val[1])

    raise ValueError('Invalid value to create Vector from: {}'.format(val))