PDFTemplateResponse and PDFTemplateView now match Django's implementations

PDFTemplateResponse is like TemplateResponse in that it does dynamic
rendering of a template on the fly.

PDFTemplateView has a much smaller implementation, relying on
PDFTemplateResponse to do the rendering for it. It also knows about
the standard TemplateResponse when it needs to render the HTML version.

1 files changed, 139 insertions, 35 deletions

diff --git a/wkhtmltopdf/views.py b/wkhtmltopdf/views.py
index 539a669..0e91b2a 100644
--- a/wkhtmltopdf/views.py
+++ b/wkhtmltopdf/views.py
@@ -2,28 +2,40 @@ from __future__ import absolute_import
 
 import os
 from re import compile
+from tempfile import NamedTemporaryFile
 import warnings
 
 from django.conf import settings
 from django.contrib.sites.models import Site
+from django.http import HttpResponse
 from django.template.context import RequestContext
-from django.template.response import HttpResponse
+from django.template.response import TemplateResponse
 from django.views.generic import TemplateView
 
-from .utils import (content_disposition_filename,
-                    template_to_temp_file, wkhtmltopdf)
+from .utils import (content_disposition_filename, wkhtmltopdf)
 
 
 class PDFResponse(HttpResponse):
     def __init__(self, content, mimetype=None, status=200,
-                 content_type='application/pdf', *args, **kwargs):
-        filename = kwargs.pop('filename', None)
-        super(PDFResponse, self).__init__(content, mimetype, status,
-                                          content_type, *args, **kwargs)
+                 content_type=None, filename=None, *args, **kwargs):
+
+        if content_type is None:
+            content_type = 'application/pdf'
+
+        super(PDFResponse, self).__init__(content=content,
+                                          mimetype=mimetype,
+                                          status=status,
+                                          content_type=content_type)
+        self.set_filename(filename)
+
+    def set_filename(self, filename):
+        self.filename = filename
         if filename:
             filename = content_disposition_filename(filename)
             header_content = 'attachment; filename={0}'.format(filename)
             self['Content-Disposition'] = header_content
+        else:
+            del self['Content-Disposition']
 
 
 class PdfResponse(PDFResponse):
@@ -33,6 +45,94 @@ class PdfResponse(PDFResponse):
         super(PdfResponse, self).__init__(content, filename=filename)
 
 
+class PDFTemplateResponse(TemplateResponse, PDFResponse):
+    """Renders a Template into a PDF using wkhtmltopdf"""
+
+    def __init__(self, request, template, context=None, mimetype=None,
+                 status=None, content_type=None, current_app=None,
+                 filename=None, header_template=None, footer_template=None,
+                 cmd_options=None, *args, **kwargs):
+
+        super(PDFTemplateResponse, self).__init__(request=request,
+                                                  template=template,
+                                                  context=context,
+                                                  mimetype=mimetype,
+                                                  status=status,
+                                                  content_type=content_type,
+                                                  current_app=None,
+                                                  *args, **kwargs)
+        self.set_filename(filename)
+
+        self.header_template = header_template
+        self.footer_template = footer_template
+
+        if cmd_options is None:
+            cmd_options = {}
+        self.cmd_options = cmd_options
+
+    def render_to_temporary_file(self, template_name, mode='w+b', bufsize=-1,
+                                 suffix='', prefix='tmp', dir=None,
+                                 delete=True):
+        template = self.resolve_template(template_name)
+        context = self.resolve_context(self.context_data)
+        content = template.render(context)
+        tempfile = NamedTemporaryFile(mode=mode, bufsize=bufsize,
+                                      suffix=suffix, prefix=prefix,
+                                      dir=dir, delete=delete)
+        try:
+            tempfile.write(content)
+            tempfile.flush()
+            return tempfile
+        except:
+            # Clean-up tempfile if an Exception is raised.
+            tempfile.close()
+            raise
+
+    @property
+    def rendered_content(self):
+        """Returns the freshly rendered content for the template and context
+        described by the PDFResponse.
+
+        This *does not* set the final content of the response. To set the
+        response content, you must either call render(), or set the
+        content explicitly using the value of this property.
+        """
+        debug = getattr(settings, 'WKHTMLTOPDF_DEBUG', False)
+
+        cmd_options = self.cmd_options.copy()
+
+        input_file = header_file = footer_file = None
+
+        try:
+            input_file = self.render_to_temporary_file(
+                template_name=self.template_name,
+                prefix='wkhtmltopdf', suffix='.html',
+                delete=(not debug)
+            )
+
+            if self.header_template:
+                header_file = self.render_to_temporary_file(
+                    template_name=self.header_template,
+                    prefix='wkhtmltopdf', suffix='.html',
+                    delete=(not debug)
+                )
+                cmd_options.setdefault('header_html', header_file.name)
+
+            if self.footer_template:
+                footer_file = self.render_to_temporary_file(
+                    template_name=self.footer_template,
+                    prefix='wkhtmltopdf', suffix='.html',
+                    delete=(not debug)
+                )
+                cmd_options.setdefault('footer_html', footer_file.name)
+
+            return wkhtmltopdf(pages=[input_file.name], **cmd_options)
+        finally:
+            # Clean up temporary files
+            for f in filter(None, (input_file, header_file, footer_file)):
+                f.close()
+
+
 class PDFTemplateView(TemplateView):
     filename = 'rendered_pdf.pdf'
     footer_template = None
@@ -42,28 +142,20 @@ class PDFTemplateView(TemplateView):
     margin_left = 0
     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':
-            return super(PDFTemplateView, self).get(request, *args, **kwargs)
-
-        if context_instance:
-            self.context_instance = context_instance
-        else:
-            self.context_instance = RequestContext(request, self.get_context_data(**kwargs))
-
-        page_path = template_to_temp_file(self.get_template_names(), self.get_context_data(), self.context_instance)
-        pdf_kwargs = self.get_pdf_kwargs()
-        output = wkhtmltopdf(page_path, **pdf_kwargs)
-        if self._tmp_files:
-            map(os.remove, self._tmp_files)
-        return self.response(output, filename=self.get_filename())
+    response_class = PDFTemplateResponse
+    html_response_class = TemplateResponse
+
+    def get(self, request, *args, **kwargs):
+        response_class = self.response_class
+        try:
+            if request.GET.get('as', '') == 'html':
+                # Use the html_response_class if HTML was requested.
+                self.response_class = self.html_response_class
+            return super(PDFTemplateView, self).get(request,
+                                                    *args, **kwargs)
+        finally:
+            # Remove self.response_class
+            self.response_class = response_class
 
     def get_filename(self):
         return self.filename
@@ -76,12 +168,6 @@ class PDFTemplateView(TemplateView):
             'margin_top': self.margin_top,
             'orientation': self.orientation,
         }
-        if self.header_template:
-            kwargs['header_html'] = template_to_temp_file(self.header_template, self.get_context_data(), self.context_instance)
-            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)
-            self._tmp_files.append(kwargs['footer_html'])
         return kwargs
 
     def get_context_data(self, **kwargs):
@@ -95,6 +181,24 @@ class PDFTemplateView(TemplateView):
 
         return context
 
+    def render_to_response(self, context, **response_kwargs):
+        """
+        Returns a PDF response with a template rendered with the given context.
+        """
+        filename = response_kwargs.pop('filename', self.get_filename())
+        cmd_options = response_kwargs.pop('cmd_options', self.get_pdf_kwargs())
+
+        if issubclass(self.response_class, PDFTemplateResponse):
+            return super(PDFTemplateView, self).render_to_response(
+                context=context, filename=filename, cmd_options=cmd_options,
+                **response_kwargs
+            )
+        else:
+            return super(PDFTemplateView, self).render_to_response(
+                context=context,
+                **response_kwargs
+            )
+
 
 class PdfTemplateView(PDFTemplateView): #TODO: Remove this in v1.0
     def __init__(self, *args, **kwargs):