calculate spanned column widths

#!/usr/bin/env python3

##  calculate spanned column widths;
##
##  ##  run
##
##  ./colwidth.py [tight|loose|empty] [left|right] < {input.txt}
##
##  -   tight: print table with tight tabskips;
##  -   loose: print table with loose tabskips;
##  -   empty: print table with empty tabskips;
##  -   left: cell text is left justified;
##  -   right: cell text is right justified;
##
##  ##  input format
##
##  input format is similar to simplified markdown table format:
##
##  -   use `|` as tabskip;
##
##  -   use `*` after `|` to indicate a col span; the number of cols spanned is
##      1 plus the number of `*`;
##
##      -   `|` is a 1-col span (namely, a normal cell);
##      -   `|*` is a 2-col span;
##      -   `|**` is a 3-col span;
##      -   ...;
##
##  -   cell text must not contain `|` or `*`;
##
##  -   within a cell, spaces surrounding text are ignored;
##
##  -   within a cell, spaces within text are kept;
##
##  ##  input example
##
##  these inputs are equivalent:
##
##      |*Ax xZ|AxxxxZ|**Axx xxx   xZ|AxxxxZ|
##      |*  Ax xZ|  AxxxxZ|**Axx xxx   xZ|AxxxxZ|
##      |*Ax xZ   |  AxxxxZ|**   Axx xxx   xZ|   AxxxxZ |

from collections import Counter
from pprint import pprint
import numpy as np
import sys

##  const: inf;
INF = 1000000

class Table:

    '''
    a table is implemented as a list of rows, with some aux data; each row is a
    list of cells; each cell is a list `[beg, end, str]`; `beg` and `end` are
    beginning and ending col indexes of this cell, inclusive; `str` is the cell
    text; rows and cells are both in sorted order;
    '''

    def __init__(self):

        ##  a list of rows;
        self.rows = []
        ##  number of rows;
        self.R = 0
        ##  number of cols;
        self.C = 0
        ##  col span counter; keyed by col index; counts how many times a cell
        ##  beginning at this column spans to or beyond the next column; for
        ##  example, if a cell has col indexes [4, 6], then increase the count
        ##  for cols 4 and 5;
        self.cnt = Counter()

    def add_row(self, line):

        '''
        parse the input line (given as str) into a row and add it to this table;
        all rows must be added before rendering this table;
        '''

        ##  row data;
        row = []

        ##  char indexes: begin and end, inclusive;
        l = INF
        r = -INF

        ##  cell indexes: begin and end, inclusive;
        i = -1
        j = -1

        ##  iter each char in input line;
        for k, c in enumerate(line):
            if c == ' ':        ##  case: ignore surrounding space;
                pass
            elif c == '|':      ##  case: works like `&` in tex;
                if i != -1:
                    ##  trick: this line works even if cell is empty;
                    row.append([i, j, line[l:r + 1]])
                    l = INF
                    r = -INF
                j += 1
                i = j
            elif c == '*':      ##  case: works like `\span` in tex;
                self.cnt[j] += 1
                j += 1
            else:               ##  case: normal char;
                l = min(l, k)
                r = max(r, k)

        ##  add this row to table;
        assert row
        self.rows.append(row)

        ##  adjust number of cols;
        if self.C == 0:
            self.C = j
        else:
            assert self.C == j

        ##  adjust number of rows;
        self.R += 1

    def merge(self):

        '''
        merge cols; a col is merged with its next col iff its col span counter
        equals to `R`, the number of rows, which means this col always spans to
        the next col;
        '''

        ##  col index mapping;
        mapping = []
        ##  col index offset;
        k = 0

        ##  calc col index mapping;
        for j in range(self.C):
            mapping.append(j - k)
            if self.cnt[j] == self.R:
                k += 1

        ##  adjust col indexes in cells;
        for row in self.rows:
            for cell in row:
                cell[0] = mapping[cell[0]]
                cell[1] = mapping[cell[1]]

        ##  adjust number of cols;
        self.C -= k

    def calc_widths(self, sep):

        '''
        calculate col widths;
        '''

        ##  maximum widths between two cols; init to zeros;
        weights = np.zeros((self.C, self.C), dtype=int)

        ##  iter all cells;
        for row in self.rows:
            for cell in row:
                i, j, s = cell[0], cell[1], cell[2]
                weights[i, j] = max(weights[i, j], len(s))

        ##  widths;
        widths = np.zeros(self.C, dtype=int)
        for j in range(self.C):
            tmp = weights[j, j]
            tot = 0
            for i in range(j - 1, -1, -1):
                tot += widths[i] + len(sep)
                tmp = max(tmp, weights[i, j] - tot)
            widths[j] = tmp

        return widths

    def pprint(self, style='loose', justify='left'):

        '''
        pretty print this table;
        '''

        ##  get left/mid/right cell separators;
        if style == 'loose':
            lsep, sep, rsep = '| ', ' | ', ' |'
        elif style == 'tight':
            lsep, sep, rsep = '|', '|', '|'
        elif style == 'empty':
            lsep, sep, rsep = '', '', ''
        else:
            raise

        ##  calculate col widths;
        widths = self.calc_widths(sep)
        print(widths)

        ##  total width;
        tot = sum(widths) + (self.C - 1) * len(sep) + len(lsep) + len(rsep)

        ##  pretty print all rows;
        for row in self.rows:
            print('-' * tot)
            print(lsep, end='')
            for k, cell in enumerate(row):
                i, j, s = cell[0], cell[1], cell[2]
                w = sum(widths[i:j + 1]) + len(sep) * (j - i)
                if justify == 'left':
                    print(s + ' ' * (w - len(s)), end='')
                elif justify == 'right':
                    print(' ' * (w - len(s)) + s, end='')
                else:
                    raise
                if k != len(row) - 1:
                    print(sep, end='')
                else:
                    print(rsep)
        print('-' * tot)

def read_table():

    '''
    read table from input;
    '''

    table = Table()
    while True:
        try:
            table.add_row(input())
        except EOFError:
            break
    return table

def main():

    '''
    main function;
    '''

    if 'empty' in sys.argv:
        style = 'empty'
    elif 'loose' in sys.argv:
        style = 'loose'
    else:
        style = 'tight'

    if 'right' in sys.argv:
        justify = 'right'
    else:
        justify = 'left'

    table = read_table()

#    print('R = {}'.format(table.R))
#    print('C = {}'.format(table.C))
#    pprint(table.rows)

    table.merge()

#    print('R = {}'.format(table.R))
#    print('C = {}'.format(table.C))
#    pprint(table.rows)

    table.pprint(style, justify)

if __name__ == '__main__':
    main()