petlutils.py

import petl
import decimal

from beancount.core import data


NUMERIC = [int, float, decimal.Decimal, data.Amount]


def render(table, colsep=' ', narrow=True, numeric=None):
    if numeric is None:
        numeric = tuple(NUMERIC)

    it = iter(table)

    # fields representation
    hdr = next(it)
    hdrrepr = list(map(str, hdr))

    # rows representations
    rows = list(it)
    rowsrepr = [list(map(str, row)) for row in rows]

    # petl supports table with uneven row lenghts
    maxrowlen = max(len(hdrrepr), max(map(len, rowsrepr)))

    # pad short rows representations to the maximum row lenght
    if len(hdrrepr) < maxrowlen:
        hdrrepr.extend([''] * (maxrowlen - len(hdrrepr)))
    for valsrepr in rowsrepr:
        if len(valsrepr) < maxrowlen:
            valsrepr.extend([''] * (maxrowlen - len(valsrepr)))

    # find longest representations so we know how wide to make cells
    colwidths = [1] * len(hdr)
    if not narrow:
        for i, hr in enumerate(hdrrepr):
            if len(hr) > 1:
                colwidths[i] = len(hr)
    for valsrepr in rowsrepr:
        for i, vr in enumerate(valsrepr):
            if len(vr) > colwidths[i]:
                colwidths[i] = len(vr)

    # construct a line for the header row
    hdrline = colsep.join(r[:w].center(w) for r, w in zip(hdrrepr, colwidths))
    sepline = colsep.join('-' * w for w in colwidths)

    def numericp(v):
        return isinstance(v, numeric) and not isinstance(v, bool)

    # construct a line for each data row
    output = [hdrline, sepline]
    for vals, valsrepr in zip(rows, rowsrepr):
        line = colsep.join(r.rjust(w) if numericp(v) else r.ljust(w)
                           for v, r, w in zip(vals, valsrepr, colwidths))
        output.append(line)

    return '\n'.join(output)


petl.Table.__str__ = render

test_petlutils.py

import decimal
import unittest
import textwrap

import petl
import petlutils

class TestPetlutils(unittest.TestCase):

    def test_render_strings(self):
        table = petl.wrap([['foo', 'bar'], ['a', 'b'], ['aaa', 'bbb']])
        r = petlutils.render(table)
        self.assertMultiLineEqual(r, textwrap.dedent("""\
          foo bar
          --- ---
          a   b  
          aaa bbb"""))

    def test_render_numeric(self):
        table = petl.wrap([['foo', 'bar'], [0, 0], [100, 1000]])
        r = petlutils.render(table)
        self.assertMultiLineEqual(r, textwrap.dedent("""\
          foo bar 
          --- ----
            0    0
          100 1000"""))

    def test_render_decimal(self):
        table = petl.wrap([['foo', 'bar'], [0, decimal.Decimal('0')], [100, decimal.Decimal('1000')]])
        r = petlutils.render(table)
        self.assertMultiLineEqual(r, textwrap.dedent("""\
          foo bar 
          --- ----
            0    0
          100 1000"""))

    def test_render_narrow_hdr(self):
        table = petl.wrap([['label'], [0], [1]])
        r = petlutils.render(table)
        self.assertMultiLineEqual(r, textwrap.dedent("""\
          l
          -
          0
          1"""))        
        
    def test_render_wide_hdr(self):
        table = petl.wrap([['label'], [0], [1]])
        r = petlutils.render(table, narrow=False)
        self.assertMultiLineEqual(r, textwrap.dedent("""\
          label
          -----
              0
              1"""))