bpeterso2000
2/15/2015 - 9:02 AM

Takes items and puts them into columns. Number of columns is automatically determined by max_width. Columns can be enumerated, colored and

Takes items and puts them into columns. Number of columns is automatically determined by max_width. Columns can be enumerated, colored and aligned (with optional fill characters). Item & columns separators can also be defined.

from toolz import partition_all

from align import align_columns
from colorize import add_colors, color_it


def get_rows_per_col(rows, num_cols):
    rows_per_col, extra = divmod(len(rows), num_cols)
    rows_per_col += bool(extra)
    return rows_per_col


def get_num_cols(rows, widths, max_width, sep):
    sep_width = len(sep)
    for num_cols, _ in enumerate(widths, 1):
        rows_per_col = get_rows_per_col(rows, num_cols)
        width = sum(map(max, partition_all(rows_per_col, widths)))
        if width + sep_width * num_cols > max_width:
            return max(num_cols - 1, 1)
    return num_cols


def number(items):
    return tuple((str(n), i) for n, i in enumerate(items, 1))


def as_columns(items, max_width=76, num_cols=-1, colors='.w.w', enum=True,
               item_sep_color=None, col_sep_color=None,
               item_sep=' ', col_sep=' | ', **alignkwds):
    """
    :param item: (2d Sequence) table to columnize
    :param max_width: (int) console width, used to determines number of columns
    :param color: (str) each letter represents a color, '.' toggles bold on/off
    :param separators: separators
    :param **alignkwds: alignment keywords
    """
    # import pdb; pdb.set_trace()
    rows = number(items) if enum else items[:]
    width_of_seps = len(item_sep) * len(rows[0][0])
    item_widths = tuple(sum(map(len, r), width_of_seps) for r in rows)
    if colors:
        rows = add_colors(rows, colors)
    if item_sep_color:
        item_sep = color_it(item_sep_color, item_sep)
    if col_sep_color:
        col_sep = color_it(col_sep_color, col_sep)
    rows = tuple(rows)
    if num_cols < 1:
        num_cols = get_num_cols(rows, item_widths, max_width, col_sep)
    if num_cols > 1:
        rows_per_col = get_rows_per_col(rows, num_cols)
        cols = (partition_all(rows_per_col, rows))
        cols = (tuple(item_sep.join(r) for r in align_columns(c, **alignkwds))
                for c in cols)
        for row in zip(*cols):
            yield col_sep.join(row)
    else:
        for row in rows:
            yield item_sep.join(tuple(align_columns([row], **alignkwds))[0])