concatenate() for numpypy

try:
    import numpypy as N
except:
    import numpy as N

TypePriorities = {
                  N.uint8:0,
                  N.uint16:1,
                  N.uint32:2,
                  N.unsignedinteger:3,
                  N.uint64:4,

                  N.int8:10,
                  N.int16:11,
                  N.int32:12,
                  int:13,
                  N.int64:14,

                  N.float32:20,
                  N.float64:21,

                  # TODO: add complex types when they will be implemented

                  object: N.inf
                  }

if 'float128' in N.__dict__:
    TypePriorities[N.float128] = 22

if 'complex' in N.__dict__:
    assert 0, 'you should update TypePriorities with complex types'

reversedTypePriorities = dict((v,k) for k, v in TypePriorities.items())

def concatenate(arrays, axis=0):
    '''
    concatenate(...)
    concatenate((a1, a2, ...), axis=0)

    Join a sequence of arrays together.

    Parameters
    ----------
    a1, a2, ... : sequence of array_like
        The arrays must have the same shape, except in the dimension
        corresponding to `axis` (the first, by default).
    axis : int, optional
        The axis along which the arrays will be joined.  Default is 0.

    Returns
    -------
    res : ndarray
        The concatenated array.
    '''
    Arrays = [(arr if isinstance(arr, N.ndarray) else N.array(arr)) for arr in arrays]
    Ndims = [arr.ndim for arr in Arrays]
    assert all([Ndims[i] == Ndims[0] for i in range(1, len(Ndims))]), 'arrays must be of same dimension'
    ndim = Ndims[0]

    #cpython numpy-style: if ndim == 1: axis = 0
    # our approach:
    assert axis < ndim, 'incorrect axis: must be less than array.ndim'

    # TODO: in numpy concatenation with numbers  or zero-shape ndarrays gives error,
    # maybe it's better to convert them automatically to ndarrays of correct shape (if possible)?

    TypePriority = TypePriorities.get(Arrays[0].dtype.type, N.inf)
    for arr in Arrays[1:]:
        if arr.ndim != ndim:
            raise ValueError('arrays must have same number of dimensions')
        tmp = TypePriorities.get(arr.dtype.type, N.inf)
        if tmp > TypePriority: TypePriority = tmp # I think it should work faster than TypePriority = max(TypePriority, tmp)
    dtype = reversedTypePriorities[TypePriority]

    Shapes = [arr.shape for arr in Arrays]
    S = Shapes[0]
    result_length_along_involved_axis = Arrays[0].shape[axis]
    ArrConcatAxisLengths = [result_length_along_involved_axis]

    for s in Shapes[1:]:
        if s[:axis] != S[:axis] or s[axis+1:] != S[axis+1:]:
            raise ValueError('array dimensions must agree except for d_0')
        ArrConcatAxisLengths.append(s[axis])
        result_length_along_involved_axis += s[axis]

    S = list(S)
    S[axis] = result_length_along_involved_axis
    r = N.empty(S, dtype)
    S = [slice(None)]*axis
    ind = 0
    for i, arr in enumerate(Arrays):
        tmp = r[S + [slice(ind, ind+ArrConcatAxisLengths[i])]]
        tmp[:] = N.array(arr, dtype) if arr.dtype != dtype else arr.copy()
        ind += ArrConcatAxisLengths[i]

    return r


# tests
assert N.all(concatenate(((1, 2), (3, 4))) == N.array((1, 2, 3, 4)))
assert N.all(concatenate((N.array((1.0, 2.0)).reshape(-1, 1), N.array((3, 4)).reshape(-1, 1)), axis=1) ==N.array(((1, 3), (2, 4))))
a = N.array([[1, 2], [3, 4]])
b = N.array([[5, 6]])
assert N.all(concatenate((a, b), axis=0) == N.array([[1, 2], [3, 4], [5, 6]]))
assert N.all(concatenate((a, b.T), axis=1) == N.array([[1, 2, 5], [3, 4, 6]]))
print('passed')