Solution to Spiral Matrix challenge with Python Generator

spiral_matrix.py

# Spiral Matrix
# https://leetcode.com/problems/spiral-matrix/solutions/4437968/solution-with-python-generator/

from typing import List


def row_seq_generator(m):
    for i in range(m):
        yield i
        if i == m - 1 - i:
            return
        yield m - 1 - i
        if 2 * (i+1) == m:
            return


def column_seq_generator(n):
    for i in range(n):
        yield n-1-i
        if i == n - 1 - i:
            return
        yield i
        if 2 * (i+1) == n:
            return


def main(matrix: List[List[int]]) -> List[int]:
    m = len(matrix)
    n = len(matrix[0])
    result = list()
    row_gen = row_seq_generator(m)
    column_gen = column_seq_generator(n)
    flag = False
    row = next(row_gen)
    column = next(column_gen)
    first_element = (0, 0)
    while True:
        try:
            second_element = (row, column)
            if first_element[0] == second_element[0]:
                step = 1
                if second_element[1] < first_element[1]:
                    step = -1
                for x in range(first_element[1], second_element[1], step):
                    result.append(matrix[first_element[0]][x])
            else:
                step = 1
                if second_element[0] < first_element[0]:
                    step = -1
                for x in range(first_element[0], second_element[0], step):
                    result.append(matrix[x][first_element[1]])
            if flag:
                column = next(column_gen)
            else:
                row = next(row_gen)
            first_element = (second_element[0], second_element[1])
            flag = not flag
        except StopIteration as e:
            i, j = second_element
            result.append(matrix[i][j])
            break
    return result


if __name__ == '__main__':
    sample_matrix = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]
    spiraled_matrix = main(sample_matrix)
    print(spiraled_matrix)