1D array library test (for Pythonista)

array.py

from __future__ import division
import numbers
import operator
import math


class array(list):
    def __getitem__(self, y):
        if isinstance(y, (numbers.Integral, slice)):
            return list.__getitem__(self, y)

        if isinstance(y[0], bool):
            inds = (i for i, yi in enumerate(y) if yi)

            return array(list.__getitem__(self, i) for i in inds)

        if isinstance(y[0], numbers.Integral):
            return array(list.__getitem__(self, yi) for yi in y)

        raise TypeError

    def __setitem__(self, keys, y):
        if isinstance(keys, (numbers.Integral, slice)):
            return list.__setitem__(self, key, y)

        if isinstance(keys[0], bool):
            inds = (i for i, ki in enumerate(keys) if ki)

            if isinstance(y, numbers.Integral):
                for i in inds:
                    list.__setitem__(self, i, y)
            else:
                for i, yi in zip(inds, y):
                    list.__setitem__(self, i, yi)

            return

        if isinstance(keys[0], numbers.Integral):
            if isinstance(y, numbers.Integral):
                for i in keys:
                    list.__setitem__(self, i, y)
            else:
                for i, yi in zip(keys, y):
                    list.__setitem__(self, i, yi)

            return

        raise TypeError

    def __abs__(self):
        return array(abs(xi) for xi in self)


def elementwise_op(op_name, right=False):
    ## get standard operator
    op = getattr(operator, '__' + op_name + '__')

    ## element-wise replacement operator
    # left operators
    def new_l_op(self_, y):
        #print 'LOP'

        if isinstance(y, numbers.Number):
            return array(op(xi, y) for xi in self_)
        else:
            return array(op(xi, yi) for xi, yi in zip(self_, y))

    # right operators
    def new_r_op(self_, y):
        #print 'ROP'

        if isinstance(y, numbers.Number):
            return array(op(y, xi) for xi in self_)
        else:
            return array(op(yi, xi) for xi, yi in zip(self_, y))

    ## replace standard operator
    if right:
        setattr(array, '__r' + op_name + '__', new_r_op)
    else:
        setattr(array, '__' + op_name + '__', new_l_op)


def patch_array():
    REPLACED_OPS = ['add', 'sub', 'mul', 'div', 'gt', 'ge', 'lt', 'le', 'pow', \
                    'or', 'and', 'xor', 'truediv', 'mod']
    REPLACED_OPS2 = ['iadd', 'isub', 'imul', 'idiv', 'ipow', 'ior', 'iand', \
                     'ixor', 'itruediv', 'imod']

    for op_name in REPLACED_OPS:
        elementwise_op(op_name)
        elementwise_op(op_name, right=True)

    for op_name in REPLACED_OPS2:
        elementwise_op(op_name)

patch_array()


def arange(*args, **kwargs):
    return array(xrange(*args, **kwargs))

def linspace(start, stop, num=50, endpoint=True, retstep=False):
    if num == 1:
        return array([start])
    elif num == 0:
        return array([])

    end_factor = 1 if endpoint else 0
    step = (stop - start)/(num - end_factor)
    space = arange(num)*step + start

    if retstep:
        return space, step

    return space

def zeros(num):
    return array([0]*num)

def ones(num):
    return array([1]*num)

def dot(x, y):
	return sum(xi*yi for xi, yi in zip(x, y))


def vectorize(func):
    if not hasattr(func, '__call__'):
        return func

    def new_func(x, *args, **kwargs):
        y = [func(xi, *args, **kwargs) for xi in x]

        return array(y)

    return new_func

def vectorize_all(namespace):
    for func in namespace.__dict__.keys():
        namespace.__dict__[func] = vectorize(namespace.__dict__[func])


vectorize_all(math)
from math import *

def main():
    x = arange(20)

    print x[(x >= 5) & (x <= 15)]

    x[x > 5] += 4
    print x[[6,7,11]]

    print sin(linspace(0, 2*pi, 5))
    print linspace(0, 7, 5, retstep=True)


if __name__ == '__main__':
    main()