Check Identity Matrix in Python

def is_identity_matrix(mat: list[list[int]]):
    n = len(mat)
    if n != len(mat[0]):
        return False

    for r in range(n):
        for c in range(n):
            if (r != c and mat[r][c] != 0) or (r == c and mat[r][c] != 1):
                return False
    return True


def read_matrix(m: int, n: int):
    mat: list[list[int]] = []
    for r in range(m):
        row = list(map(int, input(f"Row {r:2d}: ").strip().split()))
        mat.append(row)
    return mat


def main():
    n = int(input("Enter n: "))
    a = read_matrix(n, n)
    r = is_identity_matrix(a)
    print(r)


if __name__ == "__main__":
    main()