Merge pull request #9 from ecometrica/master

Refactoring of django-wkhtmltopdf

1 files changed, 252 insertions, 59 deletions

diff --git a/wkhtmltopdf/views.py b/wkhtmltopdf/views.py
index 164f94f..fff98c8 100644
--- a/wkhtmltopdf/views.py
+++ b/wkhtmltopdf/views.py
@@ -1,66 +1,283 @@
-import os
+from __future__ import absolute_import
+
 from re import compile
+from tempfile import NamedTemporaryFile
+import warnings
 
 from django.conf import settings
-from django.contrib.sites.models import Site
-from django.template.context import RequestContext
-from django.template.response import HttpResponse
+from django.http import HttpResponse
+from django.template.response import TemplateResponse
 from django.views.generic import TemplateView
 
-from wkhtmltopdf.utils import template_to_temp_file, wkhtmltopdf
+from .utils import (content_disposition_filename, override_settings,
+                    pathname2fileurl, wkhtmltopdf)
 
 
 class PDFResponse(HttpResponse):
-    def __init__(self, content, *args, **kwargs):
-        filename = kwargs.pop('filename', None)
-        super(PDFResponse, self).__init__(content, 'application/pdf', *args, **kwargs)
+    """HttpResponse that sets the headers for PDF output."""
+
+    def __init__(self, content, mimetype=None, status=200,
+                 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.__setitem__('Content-Disposition', header_content)
+            self['Content-Disposition'] = header_content
+        else:
+            del self['Content-Disposition']
 
 
 class PdfResponse(PDFResponse):
     def __init__(self, content, filename):
-        warning = '''PdfResponse is deprecated in favour of PDFResponse. It will be removed in version 1.'''
-        raise PendingDeprecationWarning(warning)
-        super(PdfResponse, self).__init__(content, filename)
+        warnings.warn('PdfResponse is deprecated in favour of PDFResponse. It will be removed in version 1.',
+                      PendingDeprecationWarning, 2)
+        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, override_settings=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
+
+        self.override_settings = override_settings
+
+    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)
+
+        # Since many things require a sensible settings.MEDIA_URL and
+        # settings.STATIC_URL, including TEMPLATE_CONTEXT_PROCESSORS;
+        # the settings themselves need to be overridden when rendering.
+        #
+        # This allows django-wkhtmltopdf to play nicely with the
+        # staticfiles app, for instance.
+        with override_settings(**self.get_override_settings()):
+            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
+
+    def convert_to_pdf(self, filename,
+                       header_filename=None, footer_filename=None):
+        cmd_options = self.cmd_options.copy()
+        # Clobber header_html and footer_html only if filenames are
+        # provided. These keys may be in self.cmd_options as hardcoded
+        # static files.
+        if header_filename is not None:
+            cmd_options['header_html'] = header_filename
+        if footer_filename is not None:
+            cmd_options['footer_html'] = footer_filename
+        return wkhtmltopdf(pages=[filename], **cmd_options)
+
+    @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', settings.DEBUG)
+
+        input_file = header_file = footer_file = None
+        header_filename = footer_filename = 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)
+                )
+                header_filename = 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)
+                )
+                footer_filename = footer_file.name
+
+            return self.convert_to_pdf(filename=input_file.name,
+                                       header_filename=header_filename,
+                                       footer_filename=footer_filename)
+        finally:
+            # Clean up temporary files
+            for f in filter(None, (input_file, header_file, footer_file)):
+                f.close()
+
+    def get_override_settings(self):
+        """Returns a dictionary of settings to override for response_class"""
+        overrides = {
+            'MEDIA_ROOT': settings.MEDIA_ROOT,
+            'MEDIA_URL': settings.MEDIA_URL,
+            'STATIC_ROOT': settings.STATIC_ROOT,
+            'STATIC_URL': settings.STATIC_URL,
+        }
+        if self.override_settings is not None:
+            overrides.update(self.override_settings)
+
+        has_scheme = compile(r'^[^:/]+://')
+
+        # If MEDIA_URL doesn't have a scheme, we transform it into a
+        # file:// URL based on MEDIA_ROOT.
+        urls = [('MEDIA_URL', 'MEDIA_ROOT'),
+                ('STATIC_URL', 'STATIC_ROOT')]
+        for url, root in urls:
+            if not has_scheme.match(overrides[url]):
+                overrides[url] = pathname2fileurl(overrides[root])
+                if not overrides[url].endswith('/'):
+                    overrides[url] += '/'
+
+        return overrides
 
 
 class PDFTemplateView(TemplateView):
+    """Class-based view for HTML templates rendered to PDF."""
+
+    # Filename for downloaded PDF. If None, the response is inline.
     filename = 'rendered_pdf.pdf'
-    footer_template = None
+
+    # Filenames for the content, header, and footer templates.
+    template_name = None
     header_template = None
-    orientation = 'portrait'
-    margin_bottom = 0
-    margin_left = 0
-    margin_right = 0
-    margin_top = 0
-    response = PDFResponse
-    _tmp_files = None
+    footer_template = None
+
+    # TemplateResponse classes for PDF and HTML
+    response_class = PDFTemplateResponse
+    html_response_class = TemplateResponse
+
+    # Command-line options to pass to wkhtmltopdf
+    cmd_options = {
+        # 'orientation': 'portrait',
+        # 'collate': True,
+        # 'quiet': 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))
+        # Copy self.cmd_options to prevent clobbering the class-level object.
+        self.cmd_options = self.cmd_options.copy()
 
-        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())
+    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
 
+    def get_cmd_options(self):
+        return self.cmd_options
+
+    def get_pdf_kwargs(self):
+        warnings.warn('PDFTemplateView.get_pdf_kwargs() is deprecated in favour of get_cmd_options(). It will be removed in version 1.',
+                      PendingDeprecationWarning, 2)
+        return self.get_cmd_options()
+
+    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', None)
+        cmd_options = response_kwargs.pop('cmd_options', None)
+        override_settings = response_kwargs.pop('override_settings', None)
+
+        if issubclass(self.response_class, PDFTemplateResponse):
+            if filename is None:
+                filename = self.get_filename()
+
+            if cmd_options is None:
+                cmd_options = self.get_cmd_options()
+
+            return super(PDFTemplateView, self).render_to_response(
+                context=context, filename=filename,
+                header_template=self.header_template,
+                footer_template=self.footer_template,
+                cmd_options=cmd_options, override_settings=override_settings,
+                **response_kwargs
+            )
+        else:
+            return super(PDFTemplateView, self).render_to_response(
+                context=context,
+                **response_kwargs
+            )
+
+
+class PdfTemplateView(PDFTemplateView): #TODO: Remove this in v1.0
+    orientation = 'portrait'
+    margin_bottom = 0
+    margin_left = 0
+    margin_right = 0
+    margin_top = 0
+
+    def __init__(self, *args, **kwargs):
+        warnings.warn('PdfTemplateView is deprecated in favour of PDFTemplateView. It will be removed in version 1.',
+                      PendingDeprecationWarning, 2)
+        super(PdfTemplateView, self).__init__(*args, **kwargs)
+
+    def get_cmd_options(self):
+        return self.get_pdf_kwargs()
+
     def get_pdf_kwargs(self):
         kwargs = {
             'margin_bottom': self.margin_bottom,
@@ -69,28 +286,4 @@ 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):
-        context = super(PDFTemplateView, self).get_context_data(**kwargs)
-
-        match_full_url = compile(r'^https?://')
-        if not match_full_url.match(settings.STATIC_URL):
-            context['STATIC_URL'] = 'http://' + Site.objects.get_current().domain + settings.STATIC_URL
-        if not match_full_url.match(settings.MEDIA_URL):
-            context['MEDIA_URL'] = 'http://' + Site.objects.get_current().domain + settings.MEDIA_URL
-
-        return context
-
-
-class PdfTemplateView(PDFTemplateView): #TODO: Remove this in v1.0
-    def as_view(cls, **initkwargs):
-        warning = '''PdfTemplateView is deprecated in favour of PDFTemplateView. It will be removed in version 1.'''
-        raise PendingDeprecationWarning(warning)
-        return super(PdfTemplateView, cls).as_view(**initkwargs)