pyltk/pyltk/tabular.py

# -*- encoding: utf-8 -*-

import itertools

from .element import element

try:
    import pandas as pd
except:
    pd = None


class row_element(element):

    """ Class representing a row of a tabular element. """

    def __init__(
        self, columns, parent, label=None, wrapper=False, endline=False, fmt=None
    ):
        super().__init__(parent, columns, label=label, raw=(fmt is not None))
        self.wrapper = wrapper
        self.endline = endline
        self.__fmt = fmt
        if self.__fmt:
            self.raw = True

    def content(self):
        row = self.childrens
        if self.__fmt:
            row = map(self.__fmt, self.childrens)
        if self.wrapper:
            cw = 1 / len(self.childrens)
            wp = lambda c: self.wrapper(c, width=cw)
            row = map(wp, row)
        out = " & ".join(map(str, row))
        if self.endline:
            out += "\\\\"
        return out


class hline_element(row_element):
    def __init__(self, command="hline", parent=None):
        super().__init__([], parent)
        self.command = command

    def content(self):
        return "\\" + self.command


# instance of hline element
hline = hline_element()


class tabular(element):
    """ Class representing a latex tabular. """

    templates = {
        "content": """\\begin{{{name}}}{{{columns}}}
{inner}
\\end{{{name}}}"""
    }

    def __init__(
        self,
        parent=None,
        rows=[],
        columns=None,
        autowrap=False,
        label=None,
        name="tabular",
    ):
        """ Create a tabular with given rows and column specification.

        Parameters:
          - rows Rows for the tabular (list of list or list of string).
          - columns String representing columns options.
        """
        super().__init__(parent, label=label)
        self.columns = columns
        self.name = name
        self.autowrap = autowrap
        for row in rows:
            self.addrow(row)

    def addhline(self):
        """ Add a hline to the current tabular element. """
        self.addrow(hline_element(parent=self), False)

    def addrow(self, row, wrap=None, fmt=None):
        """ Add a row to the tabular.

        Parameters:
          - row Array or string (array will be joined with ' & ').
          - wrap Can contains a wrapper element (typically a resizebox) for
            the cell of the row.
        """
        if wrap is None:
            wrap = self.autowrap
        if not isinstance(row, row_element):
            row = row_element(row, fmt=fmt, wrapper=wrap, parent=self)
        if self.childrens:
            last = self.childrens[-1]
            if isinstance(last, row_element):
                last.endline = True
        self.add(row)

    def addrows(self, rows, **kargs):
        """ Add multiple rows to the tabular (see addrow). """
        for row in rows:
            self.addrow(row, **kargs)

    def content(self):
        return self.fmt().format(
            "content",
            {"name": self.name, "columns": self.columns, "inner": super().content()},
        )


class multicolumn(element):
    """ Class representing a multicolumn element. """

    templates = {"content": """\\multicolumn{{{ncolumns}}}{{{align}}}{{{content}}}"""}

    def __init__(self, ncolumns, content="", alignment="c", **kargs):
        super().__init__(**kargs)
        self.ncolumns = ncolumns
        self.inner = content
        self.alignment = alignment

    def content(self):
        if self.ncolumns == 1:
            return self.inner
        return self.fmt().format(
            "content",
            {"ncolumns": self.ncolumns, "content": self.inner, "align": self.alignment},
        )


class tabulardf(tabular):
    def __init__(self, df, header=None, multilevels=False, fmt=None, **kargs):
        if header is None:
            multilevels, header = self.create_header(df)
        if multilevels:
            ncols = len(header[-1])
            df = df.reset_index()
        else:
            ncols = len(header)
            header = [header]

        kargs.setdefault("columns", "r" * ncols)
        super().__init__(**kargs)
        self.addrows(header)
        self.addhline()
        self.addrows(df.astype(object).values, fmt=fmt)

    def create_header(self, df):
        cols = df.columns
        if not isinstance(cols, pd.MultiIndex):
            return False, list(df.columns)
        idx = df.index.names
        header = []
        for levels, labels in zip(cols.levels, cols.labels):
            tmp = [multicolumn(len(idx))]
            for i, g in itertools.groupby(labels):
                tmp.append(multicolumn(len(list(g)), levels[i]))
            header.append(tmp)
        header[-1] = idx + header[-1][1:]  # Update last row
        return True, header