from io import BytesIO
import xlwt

from kg_process import connect_datebase


def export_excel(sql, head, line):
    list_obj = connect_datebase(sql)
    """导出excel表"""
    if list_obj:
        # 创建工作簿
        ws = xlwt.Workbook(encoding='utf-8')
        w = ws.add_sheet('sheet1')
        for h in range(0, len(head)):
            w.write(0, h, str(head[h]))
        excel_row = 1
        for obj in list_obj:
            for j in range(0, len(line)):
                line_name = line[j]
                w.write(excel_row, j, obj.get(line_name))
            excel_row += 1
        output = BytesIO()
        ws.save(output)
        output.seek(0)
        return output