Merge pull request #6 from incuna/develop

Fix tmp_files & tidy up response object

1 files changed, 18 insertions, 9 deletions

diff --git a/wkhtmltopdf/views.py b/wkhtmltopdf/views.py
index b15d7be..164f94f 100644
--- a/wkhtmltopdf/views.py
+++ b/wkhtmltopdf/views.py
@@ -1,4 +1,4 @@
-from os import remove
+import os
 from re import compile
 
 from django.conf import settings
@@ -11,9 +11,12 @@ from wkhtmltopdf.utils import template_to_temp_file, wkhtmltopdf
 
 
 class PDFResponse(HttpResponse):
-    def __init__(self, content, filename):
-        super(PDFResponse, self).__init__(content, 'application/pdf')
-        self.__setitem__('Content-Disposition', 'attachment; filename=%s' % filename)
+    def __init__(self, content, *args, **kwargs):
+        filename = kwargs.pop('filename', None)
+        super(PDFResponse, self).__init__(content, 'application/pdf', *args, **kwargs)
+        if filename:
+            header_content = 'attachment; filename={0}'.format(filename)
+            self.__setitem__('Content-Disposition', header_content)
 
 
 class PdfResponse(PDFResponse):
@@ -33,6 +36,11 @@ class PDFTemplateView(TemplateView):
     margin_right = 0
     margin_top = 0
     response = PDFResponse
+    _tmp_files = None
+
+    def __init__(self, *args, **kwargs):
+        self._tmp_files = []
+        super(PDFTemplateView, self).__init__(*args, **kwargs)
 
     def get(self, request, context_instance=None, *args, **kwargs):
         if request.GET.get('as', '') == 'html':
@@ -45,7 +53,10 @@ class PDFTemplateView(TemplateView):
 
         page_path = template_to_temp_file(self.get_template_names(), self.get_context_data(), self.context_instance)
         pdf_kwargs = self.get_pdf_kwargs()
-        return self.response(wkhtmltopdf(page_path, **pdf_kwargs), self.get_filename())
+        output = wkhtmltopdf(page_path, **pdf_kwargs)
+        if self._tmp_files:
+            map(os.remove, self._tmp_files)
+        return self.response(output, filename=self.get_filename())
 
     def get_filename(self):
         return self.filename
@@ -58,14 +69,12 @@ class PDFTemplateView(TemplateView):
             'margin_top': self.margin_top,
             'orientation': self.orientation,
         }
-        tmp_files = []
         if self.header_template:
             kwargs['header_html'] = template_to_temp_file(self.header_template, self.get_context_data(), self.context_instance)
-            tmp_files.append(kwargs['header_html'])
+            self._tmp_files.append(kwargs['header_html'])
         if self.footer_template:
             kwargs['footer_html'] = template_to_temp_file(self.footer_template, self.get_context_data(), self.context_instance)
-            tmp_files.append(kwargs['footer_html'])
-        map(remove, tmp_files)
+            self._tmp_files.append(kwargs['footer_html'])
         return kwargs
 
     def get_context_data(self, **kwargs):