reST simple table formatter

rstTable.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Time-stamp: <2012-07-13 21:35:19 ycopin>
# Copyright: This document has been placed in the public domain.

"""
Table examples
--------------

A simple table example:

>>> rows = [["Calvin","boy",6],
...         ["Hobbes","tiger", None],
...         ["Slimy girls"],
...         ["Susie","girl",6]]
>>> fmt = ['%6s', '%5s', '%2d']
>>> hdr=["Name", "Type", "Age"]
>>> print rst_table(rows, fmt, header=hdr)
======  =====  ===
  Name   Type  Age
======  =====  ===
Calvin    boy    6
Hobbes  tiger    -
Slimy girls
------------------
 Susie   girl    6
======  =====  ===

A more complex table, with multiple header lines, multi-columns, and
clean syntax (i.e. ready for conversion but not optimized for plain
ASCII output):

>>> rows = [["Calvin","boy",6],
...         ["Hobbes","tiger", None],
...         [["","Good match"]],
...         ["Susie","girl",6]]
>>> fmt = ['%6s', '%5s', '%3d']
>>> hdr=[["Name", "Type", "Age"],[["","Guessed"]]]
>>> tbl = Table(rows, fmt, header=hdr)
>>> tbl.construct(ascii=False)
>>> print tbl
======  =====  ===
  Name   Type  Age
..      Guessed
======  ==========
Calvin    boy    6
Hobbes  tiger   \-
..      Good match
------  ----------
 Susie   girl    6
======  =====  ===
"""

__author__ = "Yannick Copin <yannick.copin@laposte.net>"
__version__ = "Fri Jul 13 21:30:52 2012"

import numpy as N
import matplotlib.pyplot as P

import re

class Table(object):

    def __init__(self, rows, fmt, header=None):
        """See :func:`rst_table`."""

        self.rows = rows                # Data rows
        self.fmt = fmt                  # Format

        self.ncols = len(self.fmt)
        self.nrows = len(self.rows)

        if header is not None:
            if all(isinstance(cell, basestring) for cell in header):
                # Simple header: [ "h1", "h2", ... ]
                self.header = [header]
            else:
                # Multi-row header [ ["h1","h2",...], [...]]
                self.header = header
        else:
            self.header = []

        self.set_widths()
        self.totwidth = sum(self.widths) + 2*(len(self.widths)-1) # Total width

        self.construct()                # Construct self.lines

    def __str__(self):

        return '\n'.join(self.lines)

    def set_widths(self):
        """Compute column :attr:`widths` from :attr:`format`,
        increased to hold :attr:`header` labels in any."""

        self.widths = [ int(re.search('\d+',f).group()) for f in self.fmt ]
        if len(self.header)==1 and len(self.header[0])==self.ncols:
            # For plain simple header, increase width to header width
            self.widths = [ max(w,len(str(c)))
                            for w,c in zip(self.widths,self.header[0]) ]

    def separator(self, c='=', widths=None):
        """Compute separator string, made up of characters *c*."""

        return '  '.join( c*w for w in (widths
                                        if widths is not None
                                        else self.widths) )

    def construct(self, ascii=True):
        """Construct the table, i.e. list of :attr:`lines`.

        :param ascii: optimized for plain ASCII output.
        """

        noneChar = "-" if ascii else "\-"

        sep = self.separator('=')       # Build main separator

        # Build table line by line
        self.lines = [sep]              # 1st separator
        for row in self.header:         # Header lines
            if len(row)==1:
                self.lines.extend(self.multicols(row[0], ascii=ascii))
                lastRowIsMultiCol = True
            else:
                self.lines.append(
                    '  '.join( '%*s' % (w,str(c))
                               for w,c in zip(self.widths,row) ))
                lastRowIsMultiCol = False
        if self.header:                 # Add header separator
            if lastRowIsMultiCol:
                self.lines[-1] = self.lines[-1].replace('-','=')
            else:
                self.lines.append(sep)
        for row in self.rows:           # Body lines
            if len(row)==1:             # Multi-column lines
                self.lines.extend(self.multicols(row[0], ascii=ascii))
            else:
                # Embed formatted value in string of known length
                cells = [ '%*s' % (w,f % c if c is not None else noneChar)
                          for w,f,c in zip(self.widths,self.fmt,row) ]
                self.lines.append('  '.join(cells))
        self.lines.append(sep)          # Last separator

    def multicols(self, cells, ascii=True):
        """Generate complex multi-column lines to be inserted in
        table.

        :param cells: list of strings, of length <= ncols.
        :param ascii: optimized for plain ASCII output.

        If a cell (an element of *cells*) is the empty string, it is
        merged to the previous cell.
        """

        emptyCell = '' if ascii else '..'

        if isinstance(cells, basestring): # Cells is actually a single string
            cells = [cells]

        assert 0<len(cells)<=self.ncols, \
               "Multicol cells incompatible with format."
        if len(cells) < self.ncols:
            cells.extend(['']*(self.ncols-len(cells)))

        mc = []                         # Merged cell content
        mw = []                         # Merged cell width
        for cell,width in zip(cells,self.widths):
            if cell!='':                # New (merged) cell
                mc.append(cell)
                mw.append(width)
            elif mc:                    # Merge with previous cell
                mw[-1] += width + 2     # Increasing total width
            else:                       # Create the 1st (merged) cell
                mc.append(emptyCell)    # 1st cell has to be non-empty
                mw.append(width)

        lines = ['  '.join( '%-*s' % (w,str(c)) for w,c in zip(mw,mc) ),
                 self.separator('-', widths=mw)]

        if any( len(c)>w for c,w in zip(mc,mw) ):
            import warnings
            warnings.warn("""\
Table is not wide enough to accomodate requested multicols:
%s""" % '\n'.join(lines))

        return lines

    def insert_rows(self, rows, i=1):
        """Insert list of *rows* starting line *i*. The default i=1
        insert the rows right after the 1st line."""

        if i<0:                         # Allow for negative index
            i = len(self.lines) + i

        for row in rows[::-1]: # Insert in reverted order to keep index constant
            self.lines.insert(i, row)


def rst_table(rows, fmt, header=None):
    """reST simple table formatter.

    Each row of list rows (of length *nrows*) is a list of cell
    values, to be formatted by *fmt*, a list of formats (of length
    *ncols*). Header can be set to a list of column *header*.

    :param rows: input rows
    :type rows: list of *nrows* lists of 1 or *ncols* elements
    :param fmt: input format
    :type fmt: list of *ncols* format strings
    :param header: column header
    :type header: list of *ncols* header strings
    :return: formatted reST table string

    .. Note:: Input format strings have to include an explicit size,
              e.g. '%2d'.

    A row with a single cell is considered as a full-table row (if
    cell is a string), or a multi-column syntax (if cell is a list).

    .. Note:: This function is just a wrapper to
              `str(Table(rows, fmt, header=header))`

    >>> rows = [["Calvin","boy",6],
    ...         ["Hobbes","tiger", None],
    ...         ["Slimy girls"],
    ...         ["Susie","girl",6]]
    >>> fmt = ['%6s', '%5s', '%2d']
    >>> hdr=["Name", "Type", "Age"]
    >>> print rst_table(rows, fmt, header=hdr)
    ======  =====  ===
      Name   Type  Age
    ======  =====  ===
    Calvin    boy    6
    Hobbes  tiger    -
    Slimy girl
    ------------------
     Susie   girl    6
    ======  =====  ===
    """

    return str(Table(rows, fmt, header=header))

if __name__ == '__main__':

    print "\nSimple table:\n"
    rows = [["Calvin","boy",6],
            ["Hobbes","tiger", None],
            ["Slimy girls"],
            ["Susie","girl",6]]
    fmt = ['%6s', '%5s', '%2d']
    hdr=["Name", "Type", "Age"]
    print rst_table(rows, fmt, header=hdr)

    print "\nMore complex table:\n"
    rows = [["Calvin","boy",6],
            ["Hobbes","tiger", None],
            [["","Good match"]],
            ["Susie","girl",6]]
    fmt = ['%6s', '%5s', '%3d']
    hdr=[["Name", "Type", "Age"],[["","Guessed"]]]
    tbl = Table(rows, fmt, header=hdr)
    tbl.construct(ascii=False)
    print tbl