wanzb_v3_api/src/utils/exceltool.py

#-*-coding:utf-8 -*-
import StringIO
import xlwt
import xlrd
import requests
import time
from PIL import Image

def download_remote_img(url):
    """
    """
    imgname = "/tmp/" + str(int(time.time()*1000))+".bmp"
    with open(imgname,"wb+") as f:
        res = requests.get(url)
        if res.status_code == 200:
            f.write(res.content)
    return imgname


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])): 
                if j < len(data[i])-1:
                    sheet.write(i+1,j,data[i][j],style)
                else:
                    if i < 5:
                        self.insert_imgs(sheet,data[i][j],(i,j))

        output = StringIO.StringIO()
        self.book.save(output)
        self.book.save("/tmp/demo.xls")
        return output.getvalue()

    def insert_imgs(self,sheet,imgs,cell):
        """
        """
        return 
        i,j = cell
        for index,img in enumerate(imgs):
            j = j+index
            #imgpath = download_remote_img(img)
            #sheet.insert_bitmap(imgpath,i,j,scale_x=1, scale_y=1)

    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"漏洞事件"]])