TexMatrix

tex_matrix.py

# This is working pretty nice for me, and has basically paid itself off.
# In retrospect, there are some better design layouts that would be more elegant,
# if this ever becomes more than a gist: headers and footers should be their own
# class, and there should be a generic class which sandwiches the matrix with 
# such. LatexTabularX should be two of such sandwiches. Functionality for saving
# to disk and compiling the master tex file would be nice.

import warnings

class TexMatrix(object):
    """
    This class stores entries of a matrix in string format. Methods exist for populating 
    elements, submatrices, subcolumn, and subrows of this matrix. The size of the matrix
    is dynamically calculated based on the furthest x,y position you have populated thus
    far.
    Call the method `tex` to get a string that puts '&' between columns and `\\` between
    rows.
    :param str default_entry: Default entry to use in positions where no information has
        been entered.
    """
    
    TAB_WIDTH = 4
    
    def __init__(self, default_entry=''):
        self.default_entry = default_entry
        self._data = {}
        
    @property
    def n_row(self):
        """
        Current number of rows in the matrix.
        """
        pos_list = [key[0] for key in self._data.keys()]
        return 0 if len(pos_list) == 0 else max(pos_list) + 1
    @property
    def n_col(self):
        """
        Current number of columns in the matrix.
        """
        pos_list = [key[1] for key in self._data.keys()]
        return 0 if len(pos_list) == 0 else max(pos_list) + 1
        
    def populate_cell(self, value, pos):
        """
        Populate the cell at `pos` with the given value.
        
        :param str value: Value to populate with.
        :param tuple pos: `pos=(idx_row,idx_col)`
        """
        if self._data.has_key(pos):
            warnings.warn('Position {} already taken by {}...overwriting.'.format(pos, self._data[pos]))
        self._data[pos] = value
        
    def populate_submatrix(self, matrix, start_pos):
        """
        Populate cells starting at `start_pos` with the entries belonging 
        to the given square list of lists.
        
        :param list matrix: Square list of lists of strings.
        :param tuple start_pos: `start_pos=(idx_row,idx_col)`
        """
        for idx_row, row in enumerate(matrix):
            for idx_col, value in enumerate(row):
                self.populate_cell(value, (idx_row + start_pos[0], idx_col + start_pos[1]))
                
    def populate_subrow(self, row, start_pos):
        """
        Populate cells horizontally starting at `start_pos` with the entries belonging 
        to the given list of strings
        
        :param list row: List of strings.
        :param tuple start_pos: `start_pos=(idx_row,idx_col)`
        """
        self.populate_submatrix([row], start_pos)
        
    def populate_subcol(self, col, start_pos):
        """
        Populate cells vertically starting at `start_pos` with the entries belonging 
        to the given list of strings
        
        :param list col: List of strings.
        :param tuple start_pos: `start_pos=(idx_row,idx_col)`
        """
        self.populate_submatrix([[value] for value in col], start_pos)

    @property
    def matrix(self):
        """
        Return the matrix as it currently exists. Entries which have not 
        been explicited populated are given by `self.default_entry`.
        
        :returns: A list of lists of shape `(self.n_row,self.n_col)`
        """
        matrix = []
        for idx_row in range(self.n_row):
            row = []
            for idx_col in range(self.n_col):
                pos = (idx_row, idx_col)
                value = self._data[pos] if self._data.has_key(pos) else self.default_entry
                row.append(value) 
            matrix.append(row)
        return matrix
    
    @property
    def tex_list(self):
        """
        Intermediate function of `self.tex` that returns the tex code for 
        every row of the matrix.
        
        :rtype: `list`
        """
        tex = []
        for idx_row, row in enumerate(self.matrix):
            tex += [' & '.join(row) + r' \\']
        return tex
    
    @property
    def tex(self):
        """
        Returns `self.matrix` formatted as tex code, with columns spaced with `&` 
        and rows spaced with `\\` and newlines for good measure.
        
        :rtype: `str`
        """
        tab = 0
        tex_list = []
        for line in self.tex_list:
            if line.startswith(r'\end'):
                tab -= self.TAB_WIDTH
            tex_list += [(' ' * tab) + line]
            if line.startswith(r'\begin'):
                tab += self.TAB_WIDTH
        return '\n'.join(tex_list)
      
class LatexTabularX(TexMatrix):
    """
    Subclass of `TexMatrix` that wraps it with a tabularx header and footer, and possibly further wrapped 
    by a table.
    
    :param str caption: Table caption to be used, defalut is `None`.
    :param str label: Table label to be used, defalut is `None`.
    :param bool in_table: Whether or not to wrap the tabularx environment in a table environment.
    :param bool centered: Whether to include `\centering`
    :param str table_width: Width to be passed to tabularx
    :param str column_layour: String to be passed to tabularx for the column types. A string which is too short will be wrapped.
    """
    
    def __init__(self, caption=None, label=None, in_table=True, centered=True, table_width=r'\textwidth', column_layout='l'):
        if not in_table and (caption is not None or label is not None):
            raise ValueError('You need it to be in a table if you want a caption or label.')
        self.caption = caption
        self.label = label
        self.in_table = in_table
        self.centered = centered
        self.table_width = table_width
        self._column_layout = column_layout
        self._horizontal_line_locs = []
        
        super(LatexTabularX, self).__init__()
        
    @property
    def column_layout(self):
        if len(self._column_layout) > 1:
            return self._column_layout
        else:
            return self._column_layout * self.n_col
                           
    def add_horizontal_lines(self, *idxs):
        """
        Adds horizontal line following the rows with the given indeces.
        
        :param list idxs: List of row indices
        """
        self._horizontal_line_locs += idxs
        
    @property
    def tex_header(self):
        """
        Returns a list of tex code strings that form the header.
        """
        tex = []
        if self.in_table:
            tex += [r'\begin{table*}[t]']
        if self.centered:
            tex += [r'\centering']
        tex += [r'\begin{{tabularx}}{{{0}}}{{{1}}}'.format(self.table_width, self.column_layout)]
        return tex 
    
    @property
    def tex_body(self):
        """
        Returns a list of tex code strings that form the matrix itself.
        """
        tex = []
        for idx_row, row in enumerate(self.matrix):
            tex += [' & '.join(row) + r' \\']
            if idx_row in self._horizontal_line_locs:
                tex += [r'\hline \\']
        return tex
    
    @property
    def tex_footer(self):
        """
        Returns a list of tex code strings that form the footer.
        """
        tex = [r'\end{tabularx}']
        if self.caption is not None:
            tex += [r'\caption{{{0}}}'.format(self.caption)]
        if self.label is not None:
            tex += [r'\label{{{0}}}'.format(self.label)]
        if self.in_table:
            tex += [r'\end{table*}']
        return tex
    
    @property
    def tex_list(self):
        tex = self.tex_header
        tex += self.tex_body
        tex += self.tex_footer
        return tex
      
class BoldHeadedTabularX(LatexTabularX):
    """
    Identical to `LatexTabularX` except for a convenient method to put 
    a header on each column without screwing up indexing by 1.
    """
    def set_header(self, headers, add_divider=True):
        bolded = [r'\textbf{{{0}}}'.format(header) for header in headers]
        self.populate_subrow(bolded, (-1,0))
        if add_divider:
            self.add_horizontal_lines(-1)
    
    def populate_cell(self, value, pos):
        super(BoldHeadedTabularX, self).populate_cell(value, (pos[0] + 1, pos[1]))
    def add_horizontal_lines(self, *idxs):
        new_idxs = [idx+1 for idx in idxs]
        super(BoldHeadedTabularX, self).add_horizontal_lines(*new_idxs)