Source

array-experiment / array.py

Full commit
import cffi
ffi = cffi.FFI()


typecodes = {
    'c': 'char[]',
    'b': 'signed char[]',
    'B': 'unsigned char[]',
    'u': 'wchar[]',
    'h': 'signed short[]',
    'H': 'unsigned short[]',
    'i': 'signed int[]',
    'I': 'unsigned int[]',
    'l': 'signed long[]',
    'L': 'unsigned long[]',
    'f': 'float[]',
    'd': 'double[]',
}


def alloc(typecode, size):
    try:
        ctype = typecodes[typecode]
    except KeyError:
        if not isinstance(typecode, str):
            raise TypeError('must be char, not %s' %(type(typecode).__name__,))
        if len(typecode) > 1:
            raise TypeError('must be char, not str')
        raise ValueError("bad typecode (must be c, b, B, u, h, H, i, I, l, L, f or d")
    return ffi.new(ctype, size)

missing = object()


class array(object):

    def __init__(self, typecode, initializer=missing):
        self.typecode = typecode
        self.__len = 0
        self.__capacity = 1
        self.__data = alloc(typecode, 1)
        if initializer is not missing:
            self.extend(initializer)

    @property
    def itemsize(self):
        return ffi.sizeof(self.__data)/len(self.__data)

    def _realloc(self, new_size):
        assert new_size > self.__len
        new = alloc(self.typecode, new_size)
        #XXX: memcopy
        #XXX: downsizing?
        for i in range(self.__len):
            new[i] = self.__data[i]
        self.__data = new
        self.__capacity = new_size

    def __len__(self):
        return self.__len

    def append(self, item):
        #XXX: checksize
        if self.__len == self.__capacity:
            self._realloc(self.__capacity*2)
        self.__data[self.__len] = item
        self.__len += 1

    def extend(self, items):
        if isinstance(items, array) and items.typecode!=self.typecode:
            raise TypeError
        #XXX: speedups
        for item in items:
            self.append(item)



    def __add__(self, other):
        new = array(self.typecode, self)
        new.extend(other)
        return new

    def __mul__(self, other):
        #XXX int
        copy = array(self.typecode)
        copy._realloc(self.__len*other)
        for i in range(other):
            copy.extend(self)

    def __iter__(self):
        for i in range(self.__len):
            yield self.__data[i]

    def __eq__(self, other):
        if len(self) != len(other):
            return False
        return all(a == b for a, b in zip(self, other))


    def tostring(self):
        return str(ffi.buffer(self.__data, self.itemsize*self.__len))

    def fromstring(self, data):
        #XXX: expensive
        count = len(data)/self.itemsize
        temp = alloc(self.typecode, count)
        ffi.buffer(temp)[:] = data
        for i in range(count):
            self.append(temp[i])


    def tolist(self):
        return list(self)

    def __repr__(self):
        if self.__len:
            #XXX: expensive
            return 'array(%r, %r)' % (self.typecode, list(self))
        else:
            return 'array(%r)' % (self.typecode, )

    def __imul__(self, number):
        evil = array.array(self.typecode, self) *number
        self.__data = evil.__data
        self.__len = evil.__len
        self.__capacity = evil.__capacity
        return self

    def __iadd__(self, other):
        self.extend(other)
        return self