Format 2D arrays into ascii table

ascii_table.py

from typing import Any, Sequence


def ascii_table(rows: Sequence[Sequence[Any]]) -> str:
    """
    Format the rows in an ascii table.
    First row is considered to be the header
    Numbers are right justified, everything else is left justified
    See example below:
    
    ┌──────────────────────────┬─────────────────────────┬────────────────────────────────┐
    │ header longer than cells │ header of numeric value │ header shorter than cells      │
    ├──────────────────────────┼─────────────────────────┼────────────────────────────────┤
    │ line 1                   │                    1234 │ first line longer than header  │
    │ line 2                   │                    4321 │ second line longer than header │
    │ line 3                   │                   12345 │ third line longer than header  │
    └──────────────────────────┴─────────────────────────┴────────────────────────────────┘

    :param rows: 2D array of the content to format as a table
    """

    def format_cell(cell, max_width):
        if type(cell) == int:
            just = str.rjust
        else:
            just = str.ljust
        return just(str(cell), max_width, ' ')

    columns_width = [max(len(str(rows[row_idx][col_idx]))
                         for row_idx in range(len(rows)))
                     for col_idx in range(len(rows[0]))]
    top_row = '┌' + '┬'.join('─' * (width + 2) for width in columns_width) + '┐'
    header_underline = '├' + '┼'.join('─' * (width + 2) for width in columns_width) + '┤'
    bottom_row = '└' + '┴'.join('─' * (width + 2) for width in columns_width) + '┘'
    content = ['│ ' + ' │ '.join(format_cell(cell, columns_width[idx])
                                 for idx, cell in enumerate(row)) + ' │'
               for row in rows]
    ascii_rows = [top_row, content[0], header_underline] + content[1:] + [bottom_row]

    return '\n'.join(ascii_rows)

if __name__ == '__main__':
    data = [
        ['header longer than cells', 'header of numeric value', 'header shorter than cells'],
        ['line 1', 1234, 'first line longer than header'],
        ['line 2', 4321, 'second line longer than header'],
        ['line 3', 12345, 'third line longer than header'],
    ]
    table = ascii_table(data)
    print(table)

    rows_length = [len(row) for row in table.split('\n')]
    assert all(length == rows_length[0] for length in rows_length)

    rows = [row for row in table.split('\n')]

    # testing fixed position outer border characters
    assert rows[0][0] == '┌'
    assert rows[0][-1] == '┐'
    assert rows[-1][0] == '└'
    assert rows[-1][-1] == '┘'
    assert rows[2][0] == '├'
    assert rows[2][-1] == '┤'
    assert all(row[0] == row[-1] == '│' for row in rows[1:2] + rows[3:-1])

    # using top row as reference for width
    columns_width = [len(col) for col in rows[0][1:-1].split('┬')]

    # testing width of all rows except top rows as it's the reference
    assert [len(col) for col in rows[2][1:-1].split('┼')] == columns_width
    assert all([len(col) for col in row[1:-1].split('│')] == columns_width for row in rows[1:2] + rows[3:-1])
    assert [len(col) for col in rows[-1][1:-1].split('┴')] == columns_width

    # testing justification
    assert rows[3][1:-1].split('│')[0].rstrip() == ' ' + data[1][0]
    assert rows[3][1:-1].split('│')[1].lstrip() == str(data[1][1]) + ' '