extend xlsxwriter.Worksheet() class

Question

I currently override the xlsxwriter.Workbook, called rpt.Workbook. Would like to add method to xlsxwriter.Worksheet(), but since xlsxwriter.Workbook() imports Worksheet,not sure how this can be done without major convulsions. Currently, I have to pass the worksheet object as an argument.

Seems like I need to write override methods for xlsxwriter.Workbook() to point to a custom class of xlsxwriter.Worksheet() , but can't figure out how.

Here is the current override rpt.Workbook() being used:

####################################################################
class Workbook(xlsxwriter.Workbook):
####################################################################
    """\nSpreadsheet class provides methods to build a spreadsheet.
    """

    ####################################################################
    def __init__(self,filename=None, options={}):
    ####################################################################
        try:
           filename = rpt.fname(filename)
        except FileNotFoundError as err:
            log.error(err)
        return False
        log.info("Initializing excel file " + filename)
        super().__init__(filename,options)
    ####################################################################
    def add_tab(self,name=None,data=None,header=None,
            srow=0,scol=0,autocol=True):
    ####################################################################
        """\nProvides method to add_worksheet and add_table in 1 call.

             Required Attribute args:
             name = TAB name
             header = list of header names
             data = list of lists for spreadsheet contents

             Optional Attribute args:
                 srow = starting row for table, default 0
                 scol = starting col for table, default 0
                 autocol = True/False, auto set the column sizes

             add_tab also adds the worksheet.header attribute to
             allow the set_col_by_name function to work
        """

        if not data:
            log.warning("data=[][] required")
            return None
        if not header:
            log.warning("header=[] required")
            return False
        columns = []
        for field in header:
            columns.append({ 'header' : field })
        worksheet = self.add_worksheet(name)
        worksheet.header = header
        tableinfo= {
            'data' : data,
            'columns' : columns
            }
        lastcol = scol + (len(header) - 1)
        lastrow = srow + (len(data) + 1)
        worksheet.add_table(srow,scol,lastrow,lastcol,tableinfo)
       #if autocol:
           #self.auto_set_columns(worksheet=worksheet,data=data,scol=scol)
       worksheet.freeze_panes(0,1)
       return worksheet
   ####################################################################
   def auto_set_columns(self,worksheet=None,data=None,header=None,scol=0):
   ####################################################################
        """\nDetermines the max length of each column and then set
              that column width.

             Required Attribute args:
                 worksheet = worksheet object
                 data = list of lists data
             Optional Attribute args:
                 scol = Column start
                 header = row of headers for data list of lists.
                          If header not specified, worksheet
                          must have been created with self.add_tab()
        """
        if not header and worksheet.header:
            header = worksheet.header

        ## table = [] list of lists, combine header and data
        table = []
        table.append(header)        for row in data:
            table.append(row)
        ziptable = list(zip (*table))
        colsizes = []
        for idx,val in enumerate(table[0]):
            size = max(len(s) for s in ziptable[idx])
            colnum = idx + scol
            log.debug("Setting column => {} col size => {} =>  {}".format(colnum,val,size))
            worksheet.set_column(colnum,colnum,size)

I want to add a method to xlsxwriter.Worksheet() called auto_set_columns(). Currently I have to pass the worksheet object (worksheet=worksheet) as an object to get this to work.I would like to utilize worksheet.auto_set_columns() instead. Currently auto_set_columns() is a method of rpt.Workbook.

Would like auto_set_columns() to be an extended method of xlsxwriter.Worksheet.

The script side utilization currently looks like this and works:

excelfile = nashomes + '/nas_homes.xlsx'
spreadsheet = rpt.Workbook(excelfile)
worksheet = spreadsheet.add_tab(name='Nas Homes',data=hrpt.data,header=hrpt.header)
spreadsheet.auto_set_columns(worksheet=worksheet,data=hrpt.data,scol=0)

What I desire, notice the last line changes:

excelfile = nashomes + '/nas_homes.xlsx'
spreadsheet = rpt.Workbook(excelfile)
worksheet = spreadsheet.add_tab(name='Nas Homes',data=hrpt.data,header=hrpt.header)
worksheet.auto_set_columns(data=hrpt.data,scol=0)

Goal desired here is that worksheet object (which is xlsxwriter.Worksheet() ) can have an extended "auto_set_columns" method. However, since the worksheet object is created from a add_worksheet() method in the xlsxwriter.Workbook() class, I can't figure out how to extend xlsxwriter.Worksheet() without major override methods to xlsxwriter.Workbook() also. How can I get xlsxwriter.Workbook() to reference my extended xlsxwriter.Worksheet() when Workbook.add_worksheet() creates the Worksheet object?

Answer 1

How about monkey patching the worksheet before returning it from add_tab?

First create the standalone function outside of any class definitions:

import types

def auto_set_columns(self,data=None,header=None,scol=0):
    if not header and self.header:
        header = self.header

        ## table = [] list of lists, combine header and data
        table = []
        table.append(header)        
        for row in data:
            table.append(row)
        ziptable = list(zip (*table))
        colsizes = []
        for idx,val in enumerate(table[0]):
            size = max(len(s) for s in ziptable[idx])
            colnum = idx + scol
            print "Setting column => {} col size => {} =>  {}".format(colnum,val,size)
            self.set_column(colnum,colnum,size) 

And then inside your Worksheet.add_tab function, patch in the method before returning:

    ....
    worksheet.freeze_panes(0,1)
    worksheet.auto_set_columns = types.MethodType( auto_set_columns, worksheet )
    return worksheet

Now you should be able to run:

worksheet = spreadsheet.add_tab(name='Nas Homes',data=hrpt.data,header=hrpt.header)
worksheet.auto_set_columns(data=hrpt.data,scol=0)