wanzb_v3_api/src/utils/exceltool.py

#-*-coding:utf-8 -*-
import StringIO
import xlwt
import xlrd

class ExcelTool(object):
    """
    """
    def __init__(self,filename=None):
        self.filename = filename
        self.book = xlwt.Workbook(encoding="utf-8",style_compression=0)
        self.rbook = xlrd.open_workbook(filename)
        self.sheet = self.rbook.sheet_by_index(0)

    def save_data(self,sheet,header,data):
        """
        """
        sheet = self.book.add_sheet(sheet,cell_overwrite_ok=True)
        #添加表头
        for i in range(0,len(header)):
            sheet.write(0,i,header[i])

        #添加数据
        style = xlwt.easyxf('align: wrap on')
        for i in range(0,len(data)):
            for j in range(0,len(data[i])): 
                sheet.write(i+1,j,data[i][j],style)
        output = StringIO.StringIO()
        self.book.save(output)
        self.book.save("/tmp/demo.xls")
        return output.getvalue()

    def get_data(self):
        """
        """
        dct = {}
        data = []
        for i in range(1, self.sheet.nrows):
            row_list = self.sheet.row_values(i)
            data.append(row_list)
        return data


#if __name__ == "__main__":
#    mexcel = ExcelTool("/tmp/demo.xls")
#    mexcel.save_data(u"测试",[u"事件类型"],[[u"漏洞事件"]])