getShape

import numpy as np

if 'matrix' in np.__dict__:
    matrix = np.matrix
else:
    class matrix(np.ndarray):
        pass



def getShape(arrLikeItem):
    # doesn't work with sparse matrices, but it doesn't matter -
    # asarray().shape will yield also incorrect result in the case,
    # treating sparse matrices as objects
    shape = []
    tmp = arrLikeItem
    while 1:
        if isinstance(tmp, np.ndarray):
            if tmp.shape == ():
                break
            elif tmp.shape[0] == 0:
                shape += tmp.shape
                break
            elif isinstance(tmp, matrix):
                shape += tmp.shape
                break
            else:
                shape.append(tmp.shape[0])
                tmp = tmp[0]
        elif type(tmp) in (tuple, list):
            shape.append(len(tmp))
            tmp = tmp[0]
        else:
            # a number or another item which should terminate
            break
    return tuple(shape)

# version that should work with RPython:
def getShapeRPython(arrLikeItem):
    # doesn't work with sparse matrices, but it doesn't matter -
    # asarray().shape will yield also incorrect result in the case,
    # treating sparse matrices as objects
    shape = []
    Tmp = [arrLikeItem]

    while 1:
        tmp = Tmp[-1]
        if isinstance(tmp, np.ndarray):
            if tmp.shape == ():
                break
            elif tmp.shape[0] == 0:
                shape += tmp.shape
                break
            elif isinstance(tmp, matrix):
                shape += tmp.shape
                break
            else:
                shape.append(tmp.shape[0])
                Tmp.append(tmp[0])
        elif type(tmp) in (tuple, list):
            shape.append(len(tmp))
            Tmp.append(tmp[0])
        else:
            # a number or another item which should terminate
            break
    return tuple(shape)

for a in [
          1,
          [1],
          [1, 2, 3],
          np.array((1, 2, 3)),
          np.array((1, 2, 3), object),
          np.matrix(np.array((1, 2, 3))),
          np.ones((2, 3, 4)),
          np.ones((0, 3, 4)),
          (((1, 2, 3), (4, 5, 6)), ((7, 8, 9), (10, 11, 12))),
          np.array((((1, 2, 3), (4, 5, 6)), ((7, 8, 9), (10, 11, 12)))),
          ((np.array((1, 2, 3)), np.array((4, 5, 6))), (np.array((7, 8, 9)), np.array((10, 11, 12))))
          ]:
    assert np.asarray(a).shape == getShape(a) == getShapeRPython(a)
print('passed')