diff --git a/atst/routes/task_orders/index.py b/atst/routes/task_orders/index.py index 1146a9ab..97364476 100644 --- a/atst/routes/task_orders/index.py +++ b/atst/routes/task_orders/index.py @@ -1,3 +1,4 @@ +from io import BytesIO from flask import Response from . import task_orders_bp @@ -8,7 +9,8 @@ from atst.utils.docx import Docx @task_orders_bp.route("/task_orders/download_summary/") def download_summary(task_order_id): task_order = TaskOrders.get(task_order_id) - byte_str = Docx.render(data=task_order.to_dictionary()) + byte_str = BytesIO() + Docx.render(byte_str, data=task_order.to_dictionary()) filename = "{}.docx".format(task_order.portfolio_name) return Response( byte_str, diff --git a/atst/utils/docx.py b/atst/utils/docx.py index df543948..bd4b1ddd 100644 --- a/atst/utils/docx.py +++ b/atst/utils/docx.py @@ -1,5 +1,4 @@ import os -from io import BytesIO from zipfile import ZipFile from flask import render_template, current_app as app @@ -31,15 +30,14 @@ class Docx: @classmethod def render( cls, + file_like, doc_template="docx/document.xml", file_template="docx/template.docx", **args, ): document = render_template(doc_template, **args) - byte_str = BytesIO() - docx_file = ZipFile(byte_str, mode="w") - docx_template = Docx._template(file_template) - Docx._write(docx_template, docx_file, document) - docx_file.close() - byte_str.seek(0) - return byte_str + with ZipFile(file_like, mode="w") as docx_file: + docx_template = Docx._template(file_template) + Docx._write(docx_template, docx_file, document) + file_like.seek(0) + return file_like diff --git a/tests/utils/test_docx.py b/tests/utils/test_docx.py index e35e2fc4..9b643609 100644 --- a/tests/utils/test_docx.py +++ b/tests/utils/test_docx.py @@ -6,7 +6,8 @@ from atst.utils.docx import Docx def test_render_docx(): data = {"droid_class": "R2"} - docx_file = Docx.render(data=data) + byte_str = BytesIO() + docx_file = Docx.render(byte_str, data=data) zip_ = ZipFile(docx_file, mode="r") document = zip_.read(Docx.DOCUMENT_FILE) assert b"droid_class: R2" in document