From 9aeb68c40245624f17f3d04a0457d3f953c8edb6 Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <shaohu.zhang@fit2cloud.com>
Date: Sun, 29 Sep 2024 16:02:59 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E6=AE=B5=E8=90=BD?=
 =?UTF-8?q?=E8=BF=87=E9=95=BF=E5=AF=BC=E5=87=BA=E7=9F=A5=E8=AF=86=E5=BA=93?=
 =?UTF-8?q?=E6=8A=A5=E9=94=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/dataset/serializers/dataset_serializers.py  |  2 +-
 apps/dataset/serializers/document_serializers.py | 16 +++++++++-------
 ui/src/api/dataset.ts                            |  2 +-
 ui/src/api/document.ts                           |  2 +-
 4 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/apps/dataset/serializers/dataset_serializers.py b/apps/dataset/serializers/dataset_serializers.py
index 2fdfd9b125..d48b9d8f6b 100644
--- a/apps/dataset/serializers/dataset_serializers.py
+++ b/apps/dataset/serializers/dataset_serializers.py
@@ -681,7 +681,7 @@ def export_excel(self, with_valid=True):
                                                                                  document_list)
             workbook = DocumentSerializers.Operate.get_workbook(data_dict, document_dict)
             response = HttpResponse(content_type='application/vnd.ms-excel')
-            response['Content-Disposition'] = 'attachment; filename="dataset.xls"'
+            response['Content-Disposition'] = 'attachment; filename="dataset.xlsx"'
             workbook.save(response)
             return response
 
diff --git a/apps/dataset/serializers/document_serializers.py b/apps/dataset/serializers/document_serializers.py
index edc7b6c22b..91ade28fc3 100644
--- a/apps/dataset/serializers/document_serializers.py
+++ b/apps/dataset/serializers/document_serializers.py
@@ -14,7 +14,7 @@
 from functools import reduce
 from typing import List, Dict
 
-import xlwt
+import openpyxl
 from celery_once import AlreadyQueued
 from django.core import validators
 from django.db import transaction
@@ -34,8 +34,8 @@
 from common.handle.impl.qa.xls_parse_qa_handle import XlsParseQAHandle
 from common.handle.impl.qa.xlsx_parse_qa_handle import XlsxParseQAHandle
 from common.handle.impl.table.csv_parse_table_handle import CsvSplitHandle
-from common.handle.impl.table.xlsx_parse_table_handle import XlsxSplitHandle
 from common.handle.impl.table.xls_parse_table_handle import XlsSplitHandle
+from common.handle.impl.table.xlsx_parse_table_handle import XlsxSplitHandle
 from common.handle.impl.text_split_handle import TextSplitHandle
 from common.mixins.api_mixin import ApiMixin
 from common.util.common import post, flat_map
@@ -490,25 +490,27 @@ def export(self, with_valid=True):
             data_dict, document_dict = self.merge_problem(paragraph_list, problem_mapping_list, [document])
             workbook = self.get_workbook(data_dict, document_dict)
             response = HttpResponse(content_type='application/vnd.ms-excel')
-            response['Content-Disposition'] = f'attachment; filename="data.xls"'
+            response['Content-Disposition'] = f'attachment; filename="data.xlsx"'
             workbook.save(response)
             return response
 
         @staticmethod
         def get_workbook(data_dict, document_dict):
             # 创建工作簿对象
-            workbook = xlwt.Workbook(encoding='utf-8')
+            workbook = openpyxl.Workbook()
+            workbook.remove_sheet(workbook.active)
             for sheet_id in data_dict:
                 # 添加工作表
-                worksheet = workbook.add_sheet(document_dict.get(sheet_id))
+                worksheet = workbook.create_sheet(document_dict.get(sheet_id))
                 data = [
                     ['分段标题（选填）', '分段内容（必填，问题答案，最长不超过4096个字符）', '问题（选填，单元格内一行一个）'],
-                    *data_dict.get(sheet_id)
+                    *data_dict.get(sheet_id, [])
                 ]
                 # 写入数据到工作表
                 for row_idx, row in enumerate(data):
                     for col_idx, col in enumerate(row):
-                        worksheet.write(row_idx, col_idx, col)
+                        cell = worksheet.cell(row=row_idx + 1, column=col_idx + 1)
+                        cell.value = col
                     # 创建HttpResponse对象返回Excel文件
             return workbook
 
diff --git a/ui/src/api/dataset.ts b/ui/src/api/dataset.ts
index 6e61db1df8..8b92598614 100644
--- a/ui/src/api/dataset.ts
+++ b/ui/src/api/dataset.ts
@@ -199,7 +199,7 @@ const exportDataset: (
   dataset_id: string,
   loading?: Ref<boolean>
 ) => Promise<any> = (dataset_name, dataset_id, loading) => {
-  return exportExcel(dataset_name + '.xls', `dataset/${dataset_id}/export`, undefined, loading)
+  return exportExcel(dataset_name + '.xlsx', `dataset/${dataset_id}/export`, undefined, loading)
 }
 
 export default {
diff --git a/ui/src/api/document.ts b/ui/src/api/document.ts
index d4d9b505b0..371a2d4db0 100644
--- a/ui/src/api/document.ts
+++ b/ui/src/api/document.ts
@@ -310,7 +310,7 @@ const exportDocument: (
   loading?: Ref<boolean>
 ) => Promise<any> = (document_name, dataset_id, document_id, loading) => {
   return exportExcel(
-    document_name + '.xls',
+    document_name + '.xlsx',
     `${prefix}/${dataset_id}/document/${document_id}/export`,
     {},
     loading