瀏覽代碼

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)
     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'])
 
 

+ 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.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import StringIO
 """
 Common library to export either CSV or XLS.
 """
 import gc
 import logging
 import tablib
+import xlsxwriter
 
 from django.http import StreamingHttpResponse
 from django.utils.encoding import smart_str
@@ -54,6 +56,34 @@ def dataset(headers, data, encoding=None):
   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):
   if format == 'csv':
     show_headers = True
@@ -83,7 +113,7 @@ def create_generator(content_generator, format, encoding=None):
     if len(data) > MAX_XLS_ROWS:
       data = data[:MAX_XLS_ROWS]
 
-    yield dataset(headers, data, encoding).xls
+    yield xls_dataset(headers, data, encoding).xls
     gc.collect()
 
 
@@ -97,7 +127,7 @@ def make_response(generator, format, name, encoding=None):
   if format == 'csv':
     content_type = 'application/csv'
   elif format == 'xls':
-    content_type = 'application/xls'
+    content_type = 'application/xlsx'
   elif format == 'json':
     content_type = 'application/json'
   else: