Print tabular data to a console making a best effort to write one row per line, adjusting column sizes and truncating automatically to fit each row.

console_table.py

"""
Print tabular data to a console making a best effort to write one row per line,
adjusting column sizes and truncating automatically to fit each row.

The main function is `console_table.fmt`. Execute this script to run the sample.

https://gist.github.com/dirkraft/f0759cebfd4d4567365c36570f1e3e51

Written against python 3.6.3. Requires pydash (sorry purists, I like this library :D ).


The MIT License (MIT)

Copyright (c) 2018 Jason Dunkelberger (a.k.a. "dirkraft")

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import os
import re
from itertools import cycle
from subprocess import getoutput
from typing import Any, Callable, Dict, List, Union

from pydash import py_

console_width = py_.find(py_.map([
    os.environ.get('COLUMNS'),
    getoutput('stty size').split()[-1],
    120
], py_.to_number))

markup_ascii = {
    '-': '-',
    '|': '|',
    '+': '+',
    '~': '~',
}
markup_unicode = {
    '-': '─',  # 0x2500
    '|': '│',  # 0x2502
    '+': '┼',  # 0x253c
    '~': '~',  # ascii
}
markup_default = markup_ascii

_ansi_reset = '\x1b[39m'
_ansi_re = re.compile(r'\x1b\[\d{2}m')


def distribute_fairly(width: int, col_widths: [int]) -> [int]:
    """
    Distributes width "fairly" across columns.

    This function can be used in `fmt(distributor=distribute_fairly)`
    "Fairly" means that all columns get equal *chance* at available columns.

    Example: width=40, col_widths={0: 10, 1: 20, 2: 30}
    - Column 0 gets all 10 since 10 * 3 < 40
    - The remaining columns {1: 20, 2: 30} must share 30 remaining width.
    - Remaining columns cannot be granted requested width since 30 / 2 < 20.
    - They share remaining width.
    - Fair distribution comes out to be {0: 10, 1: 15, 2: 15}

    :param width: total width available. Columns will be reduced to fit
    into this width if they cannot fit.
    :param col_widths: Desired column widths.
    :return: best possible width for each column
    """

    def recurse(w: int, cols: Dict[int, int]):
        fair_share = w // len(cols)
        fair_cols, unfair_cols = py_(cols.items()) \
            .partition(lambda col: col[1] <= fair_share) \
            .map(dict) \
            .value()

        if len(fair_cols) == len(cols):
            # Everything is fair.
            assert fair_cols and not unfair_cols
            return cols

        elif len(unfair_cols) == len(cols):
            # Nothing fits. Truncate them all.
            assert not fair_cols and unfair_cols

            truncated_width = int(w // len(unfair_cols))
            # To recover widths lost in int division remainders, add 1 to some cols.
            lost_to_int_div = w - (truncated_width * len(unfair_cols))

            return {col_idx: truncated_width + int(col_idx < lost_to_int_div)
                    for col_idx, _ in unfair_cols.items()}

        else:
            # Some cols fit in the current space fairly. Keep those and
            # divide up the rest with the remaining space.
            assert fair_cols and unfair_cols

            fair_usage = py_.sum(fair_cols)
            remaining_width = w - fair_usage

            redistributed_unfair_cols = recurse(remaining_width, unfair_cols)
            return py_.merge(fair_cols, redistributed_unfair_cols)

    resized = recurse(width, py_.to_dict(col_widths))
    return py_(resized.items()).sort_by(0).map(1).value()


def justify_fairly(width: int, col_widths: [int]) -> [int]:
    """
    Same as distribute_fairly but always fills all columns.

    Once available columns have been distribute fairly, if the table is undersized.
    Distributes one extra space to each column starting with the smallest.
    If the the table is STILL undersized repeats this.

    :param width: total width available. Columns will be reduced to fit
    into this width if they cannot fit.
    :param col_widths: Desired column widths.
    :return: best possible width for each column
    """
    fair = distribute_fairly(width, col_widths)

    rem = width - sum(fair)
    smallest_col_idx_first = py_(enumerate(fair)).sort_by(1).map(0).value()
    round_robin = cycle(smallest_col_idx_first)
    while rem > 0:
        col_idx = next(round_robin)
        fair[col_idx] += 1
        rem -= 1

    return fair


def fmt(data: [Any],
        headers: [str] = None,
        formatters: [Callable[[Any], Any]] = None,
        formatter: Callable[[Any], List[Any]] = None,
        table_width: int = None,
        table_adj: int = 0,
        markup: Dict[str, str] = None,
        new_line_repl: str = ' ',
        out: [Callable[[Any], Any]] = print,
        ansi: bool = False,
        distributor: Callable[[int, List[int]], List[int]] = distribute_fairly,
        align: Union[str, Callable[[str], Callable[[int], str]]] = None):
    """
    Formats rows of data for output to a console.

    Makes a pass over the data to distribute available width to each
    column, truncating columns which have cells that are too long if necessary.
    Each row must fit within one console line.

    :param data: Objects which form the basis of each row.
    :param headers: Optional header row to start the table.
    :param formatters: A list of functions one per column, which each
                       generate a cell of the row for each data object.
    :param formatter: A single function which produces the
                      entire row for each data object.
    :param table_width: Max table width. Defaults to `console_width`.
    :param table_adj: Manual adjustment to automatic width.
    :param markup: Markup used to render table. Keys are [|-+~] and are the defaults.
                   ~ symbolizes truncation.
    :param new_line_repl: String which replaces new lines inside of cell values.
                          Defaults to a space ' '.
    :param out: Where to send the output. Defaults to `print` on stdout.
    :param ansi: Set to True to preserve ANSI codes in output. Defaults to False.
    :param distributor: How to divide up available width across columns.
    :param align: How to align values in columns. Default left aligns strings, right aligns nums.
                  Pass 'l' for always left or 'r' for always right.
    """
    table_width = table_width or console_width

    if markup is None:
        markup = markup_default
    # Requires all data to be buffered to find optimal column sizes.
    data = list(data)
    if headers == 'keys':
        if len(data) > 0:
            headers = data[0].keys()
        else:
            headers = []
    # Recreate as list for predictability.
    headers = None if headers is None else list(headers)

    # Use headers as keys? Otherwise don't prevent indexed elements.
    if formatters is None and len(data) > 0 and hasattr(data[0], 'keys') and headers:
        formatters = headers
    # Recreate as list again for predictability.
    formatters = None if formatters is None else list(formatters)

    # First try the strong clues.
    num_cols = max(
        # headers given?
        0 if headers is None else len(headers),
        # formatters given?
        0 if formatters is None else len(formatters),
    )
    if num_cols == 0:
        # Then some weaker clues
        first_data_peek = py_.get(data, 0, [])
        num_cols = max(
            # already a row?
            len(first_data_peek),
            # formatter given and there's a first datum to row-ify?
            len(formatter(first_data_peek)) if first_data_peek and formatter else 0,
        )
    if num_cols == 0:
        # There's nothing to print.
        return ''

    def formatters_rollup(o):
        def format_o(f):
            if isinstance(f, str) or isinstance(f, int):
                formatted = py_.get(o, f)
            else:
                formatted = f(o)

            if formatted is None:
                return ""
            else:
                s = str(formatted)
                s = s.replace("\r\n", new_line_repl)
                s = s.replace("\r", new_line_repl)
                s = s.replace('\n', new_line_repl)
                return s

        return py_.map(formatters, format_o)

    if not formatter:
        if not formatters:
            if len(data) > 0 and py_.is_function(getattr(data[0], 'keys', None)):
                # Access row elements by key. Formatters are key gets.
                formatters = data[0].keys()
            else:
                # Assume list-like, access row elements by index. Formatters are just indices.
                formatters = list(range(0, num_cols))
        formatter = formatters_rollup

    width_used_by_seps = len(markup['|']) * (num_cols - 1)
    avail_width = table_width + table_adj - width_used_by_seps
    assert avail_width >= num_cols, 'Width is too small to even show one character per column.'

    data_rows = py_.map(data, formatter)
    all_rows = data_rows
    if headers is not None:
        all_rows = [headers] + data_rows

    def get_len(s: str) -> int:
        if ansi:
            stripped = re.sub(_ansi_re, '', s)
            return py_.size(stripped)
        else:
            return py_.size(s)

    def get_col_max(col_idx: int):
        return py_(all_rows).map(lambda row: py_.get(row, col_idx, '')).map(get_len).max().value()

    max_col_widths = [get_col_max(i) for i in range(0, num_cols)]
    resized_col_widths = distributor(avail_width, max_col_widths)

    def is_number(cell: str):
        if ansi:
            cell = re.sub(_ansi_re, '', cell)
        try:
            float(cell)
            return True
        except ValueError:
            return False

    def out_row(row: [str], sep: str = markup['|']):
        cells = py_.zip(row, resized_col_widths)
        line = []
        for (idx, (cell, max_width)) in enumerate(cells):
            cell_len = get_len(cell)
            if cell_len > max_width:
                cell = cell[:max_width - len(markup['~'])] + markup['~']
                # Super edge case, but do what we can even if the truncator doesn't fit.
                cell = cell[:max_width]

            if ansi and not cell.endswith(_ansi_reset) and _ansi_re.search(cell):
                cell += _ansi_reset

            if align == 'l':
                def padding_fn(s):
                    return s.ljust
            elif align == 'r':
                def padding_fn(s):
                    return s.rjust
            elif align:
                padding_fn = align
            else:
                def padding_fn(s):
                    return s.rjust if is_number(s) else s.ljust

            recover_ansi_width = len(cell) - cell_len
            padded = padding_fn(cell)(max_width + recover_ansi_width)

            line.append(padded)
            if idx < len(cells) - 1:
                line.append(sep)
        out(''.join(line))

    if headers is not None:
        out_row(headers)
        out_row([markup['-'] * width for width in resized_col_widths], sep=markup['+'])

    py_.map(data_rows, lambda row: out_row(row, sep=markup['|']))


def _fmt_test():
    fmt(
        headers=('something', 'something else', 'how much', 'heyhey'),
        data=[
            ('Antidisestablishmentarianism', 'noun', 75, 'Whatever you want it to be.'),
            ('way' + 'y' * 90 + ' too long', 'phrase', 1000000000, 'haha'),
            ('ansi_test \x1b[32mgreen', '', '', ''),
        ],
        table_width=100,
        ansi=True,
        markup=markup_unicode,
    )


if __name__ == '__main__':
    _fmt_test()