A function for diagonal matrix traversal. (python3)

diagtravmat.py

def diag_trav_mat(mat,func):
    # pretend to be working with a square matrix
    # extend length to the largest of the sides
    n = max(len(mat),len(mat[0]))
    
    # for the amount of diagnoals 
    for x in range(0, 2*n-1):
        # num of elements in diagonal
        el = n - abs(n-1-x)
        
        # row (i) range, the column (j) range is the same but reversed
        r = range(min(x, n-1), min(x,n-1)-el,-1)

        for i,j in zip(r, reversed(r)):
            try:
                # apply func to element
                func(mat[i][j])
            except IndexError:
                # in case of non-square matrices
                # we expect to get out of bounds
                pass


### TEST CASES ###
# test matrix n = 2
m2 = [[1,3],[2,4]]
# test matrix n = 3
m3 = [[1,3,6],[2,5,8],[4,7,9]]
# test matrix n = 4
m4 = [[1,3,6,10],[2,5,9,13],[4,8,12,15],[7,11,14,16]]

# non square test matrix n = 2 m = 3
mn23 = [[1,3,5],[2,4,6]]
# non square test matrix n = 2 m = 3
mn32 = [[1,3],[2,5],[4,6]]

# print elements
diag_trav_mat(m3,print)