import numpypy 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), #<-works in CPython but not PyPy
          #np.matrix(np.array((1, 2, 3))), # <-works in CPython but not PyPy - matrix is absent yet
          np.ones((2, 3, 4)),
          #np.ones((0, 3, 4)), #np.array(a).shape gives incorrect shape of a: (0,), while it is  (0,3,4) and my funcs return correct result
          (((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))))
          ]:

    # print np.array(a).shape ,  getShape(a) ,  getShapeRPython(a)
    assert np.array(a).shape == getShape(a) == getShapeRPython(a)
print('passed')