浏览代码

HUE-2523 [core] Migrate Excel generator to XlsxWriter

Romain Rigaux 10 年之前
父节点
当前提交
65b4918
共有 2 个文件被更改,包括 33 次插入3 次删除
  1. 1 1
      apps/search/src/search/tests.py
  2. 32 2
      desktop/core/src/desktop/lib/export_csvxls.py

+ 1 - 1
apps/search/src/search/tests.py

@@ -277,7 +277,7 @@ class TestWithMockedSolr(TestSearchBase):
     })
     })
     xls_response_content = ''.join(xls_response.streaming_content)
     xls_response_content = ''.join(xls_response.streaming_content)
     assert_not_equal(0, len(xls_response_content))
     assert_not_equal(0, len(xls_response_content))
-    assert_equal('application/xls', xls_response['Content-Type'])
+    assert_equal('application/xlsx', xls_response['Content-Type'])
     assert_equal('attachment; filename=query_result.xls', xls_response['Content-Disposition'])
     assert_equal('attachment; filename=query_result.xls', xls_response['Content-Disposition'])
 
 
 
 

+ 32 - 2
desktop/core/src/desktop/lib/export_csvxls.py

@@ -14,12 +14,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # limitations under the License.
+import StringIO
 """
 """
 Common library to export either CSV or XLS.
 Common library to export either CSV or XLS.
 """
 """
 import gc
 import gc
 import logging
 import logging
 import tablib
 import tablib
+import xlsxwriter
 
 
 from django.http import StreamingHttpResponse
 from django.http import StreamingHttpResponse
 from django.utils.encoding import smart_str
 from django.utils.encoding import smart_str
@@ -54,6 +56,34 @@ def dataset(headers, data, encoding=None):
   return dataset
   return dataset
 
 
 
 
+class XlsWrapper():
+  def __init__(self, xls):
+    self.xls = xls
+
+
+def xls_dataset(headers, data, encoding=None):
+  output = StringIO.StringIO()
+
+  workbook = xlsxwriter.Workbook(output)
+  worksheet = workbook.add_worksheet()
+
+  n = 0
+
+  if headers:
+    worksheet.write_row(n, 0, format(headers, encoding))
+    n +=1
+
+  for row in data:
+    worksheet.write_row(n, 0, format(row, encoding))
+    n +=1
+
+  workbook.close()
+
+  output.seek(0)
+
+  return XlsWrapper(output.read())
+
+
 def create_generator(content_generator, format, encoding=None):
 def create_generator(content_generator, format, encoding=None):
   if format == 'csv':
   if format == 'csv':
     show_headers = True
     show_headers = True
@@ -83,7 +113,7 @@ def create_generator(content_generator, format, encoding=None):
     if len(data) > MAX_XLS_ROWS:
     if len(data) > MAX_XLS_ROWS:
       data = data[:MAX_XLS_ROWS]
       data = data[:MAX_XLS_ROWS]
 
 
-    yield dataset(headers, data, encoding).xls
+    yield xls_dataset(headers, data, encoding).xls
     gc.collect()
     gc.collect()
 
 
 
 
@@ -97,7 +127,7 @@ def make_response(generator, format, name, encoding=None):
   if format == 'csv':
   if format == 'csv':
     content_type = 'application/csv'
     content_type = 'application/csv'
   elif format == 'xls':
   elif format == 'xls':
-    content_type = 'application/xls'
+    content_type = 'application/xlsx'
   elif format == 'json':
   elif format == 'json':
     content_type = 'application/json'
     content_type = 'application/json'
   else:
   else: