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.
author: Simon Law <simon.law@ecometrica.com> 2012-07-24 14:57:06 -0400
committer: Simon Law <simon.law@ecometrica.com> 2012-07-24 14:57:06 -0400
commit: 8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7 (patch)
tree: e88c55a0ba32e5222547a2edb02a589a6cad6e84 /wkhtmltopdf
parent: 4a27fe9f16ecd4bc6f94ad89046767450316490c (diff)
download: django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.tar.gz
django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.tar.bz2
django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.zip
3 files changed, 307 insertions, 38 deletions
diff --git a/wkhtmltopdf/tests.py b/wkhtmltopdf/tests.py
index 6b49378..0497dd3 100644
--- a/wkhtmltopdf/tests.py
+++ b/wkhtmltopdf/tests.py
@@ -7,10 +7,13 @@ import sys
 import warnings
 
 from django.test import TestCase
+from django.test.client import RequestFactory
 
 from .subprocess import CalledProcessError
-from .utils import _options_to_args, template_to_temp_file, wkhtmltopdf
-from .views import PDFResponse, PdfResponse, PdfTemplateView
+from .utils import (override_settings,
+                    _options_to_args, template_to_temp_file, wkhtmltopdf)
+from .views import (PDFResponse, PdfResponse, PDFTemplateResponse,
+                    PDFTemplateView, PdfTemplateView)
 
 
 class TestUtils(TestCase):
@@ -70,7 +73,7 @@ class TestUtils(TestCase):
 
 class TestViews(TestCase):
     def test_pdf_response(self):
-        """Should generate the correct HttpResonse object and mimetype"""
+        """Should generate the correct HttpResponse object and mimetype"""
         # 404
         response = PDFResponse(content='', status=404)
         self.assertEqual(response.status_code, 404)
@@ -108,6 +111,113 @@ class TestViews(TestCase):
                                mimetype='application/x-pdf')
         self.assertEqual(response['Content-Type'], 'application/x-pdf')
 
+    def test_pdf_template_response(self):
+        """Test PDFTemplateResponse."""
+        from django.conf import settings
+
+        with override_settings(
+            MEDIA_URL='/media/',
+            STATIC_URL='/static/',
+            TEMPLATE_LOADERS=['django.template.loaders.filesystem.Loader'],
+            TEMPLATE_DIRS=[os.path.join(os.path.dirname(__file__),
+                                        'testproject', 'templates')],
+            WKHTMLTOPDF_DEBUG=False,
+        ):
+            # Setup sample.html
+            template = 'sample.html'
+            context = {'title': 'Heading'}
+            request = RequestFactory().get('/')
+            response = PDFTemplateResponse(request=request,
+                                           template=template,
+                                           context=context)
+            self.assertEqual(response._request, request)
+            self.assertEqual(response.template_name, template)
+            self.assertEqual(response.context_data, context)
+            self.assertEqual(response.filename, None)
+            self.assertEqual(response.header_template, None)
+            self.assertEqual(response.footer_template, None)
+            self.assertEqual(response.cmd_options, {})
+            self.assertFalse(response.has_header('Content-Disposition'))
+
+            # Render to temporary file
+            tempfile = response.render_to_temporary_file(template)
+            tempfile.seek(0)
+            html_content = tempfile.read()
+            self.assertTrue(html_content.startswith('<html>'))
+            self.assertTrue('<h1>{title}</h1>'.format(**context)
+                            in html_content)
+
+            pdf_content = response.rendered_content
+            self.assertTrue(pdf_content.startswith('%PDF-'))
+            self.assertTrue(pdf_content.endswith('%%EOF\n'))
+
+            # Header
+            filename = 'output.pdf'
+            footer_template = 'footer.html'
+            cmd_options = {'title': 'Test PDF'}
+            response = PDFTemplateResponse(request=request,
+                                           template=template,
+                                           context=context,
+                                           filename=filename,
+                                           footer_template=footer_template,
+                                           cmd_options=cmd_options)
+            self.assertEqual(response.filename, filename)
+            self.assertEqual(response.header_template, None)
+            self.assertEqual(response.footer_template, footer_template)
+            self.assertEqual(response.cmd_options, cmd_options)
+            self.assertTrue(response.has_header('Content-Disposition'))
+
+            tempfile = response.render_to_temporary_file(footer_template)
+            tempfile.seek(0)
+            footer_content = tempfile.read()
+            self.assertTrue('MEDIA_URL = {}'.format(settings.MEDIA_URL)
+                            in footer_content)
+            self.assertTrue('STATIC_URL = {}'.format(settings.STATIC_URL)
+                            in footer_content)
+
+            pdf_content = response.rendered_content
+            self.assertTrue('\0'.join('{title}'.format(**cmd_options))
+                            in pdf_content)
+
+    def test_pdf_template_view(self):
+        """Test PDFTemplateView."""
+        with override_settings(
+            MEDIA_URL='/media/',
+            STATIC_URL='/static/',
+            TEMPLATE_LOADERS=['django.template.loaders.filesystem.Loader'],
+            TEMPLATE_DIRS=[os.path.join(os.path.dirname(__file__),
+                                        'testproject', 'templates')],
+            WKHTMLTOPDF_DEBUG=False,
+        ):
+            # Setup sample.html
+            template = 'sample.html'
+            filename = 'output.pdf'
+            view = PDFTemplateView.as_view(filename=filename,
+                                           template_name=template)
+
+            # As PDF
+            request = RequestFactory().get('/')
+            response = view(request)
+            self.assertEqual(response.status_code, 200)
+            response.render()
+            self.assertEqual(response['Content-Disposition'],
+                             'attachment; filename="{}"'.format(filename))
+            self.assertTrue(response.content.startswith('%PDF-'))
+            self.assertTrue(response.content.endswith('%%EOF\n'))
+
+            # As HTML
+            request = RequestFactory().get('/?as=html')
+            response = view(request)
+            self.assertEqual(response.status_code, 200)
+            response.render()
+            self.assertFalse(response.has_header('Content-Disposition'))
+            self.assertTrue(response.content.startswith('<html>'))
+
+            # POST
+            request = RequestFactory().post('/')
+            response = view(request)
+            self.assertEqual(response.status_code, 405)
+
     def test_deprecated(self):
         """Should warn when using deprecated views."""
         with warnings.catch_warnings(record=True) as w:
diff --git a/wkhtmltopdf/utils.py b/wkhtmltopdf/utils.py
index fb8f0aa..d0475c4 100644
--- a/wkhtmltopdf/utils.py
+++ b/wkhtmltopdf/utils.py
@@ -5,6 +5,7 @@ from itertools import chain
 from os import fdopen
 import sys
 from tempfile import mkstemp
+import warnings
 
 from django.conf import settings
 from django.template import loader
@@ -76,6 +77,8 @@ def template_to_temp_file(template_name, dictionary=None, context_instance=None)
     """
     Renders a template to a temp file, and returns the path of the file.
     """
+    warnings.warn('template_to_temp_file is deprecated in favour of PDFResponse. It will be removed in version 1.',
+                  PendingDeprecationWarning, 2)
     file_descriptor, tempfile_path = mkstemp(suffix='.html')
     with fdopen(file_descriptor, 'wt') as f:
         f.write(smart_str(loader.render_to_string(template_name, dictionary=dictionary, context_instance=context_instance)))
@@ -111,3 +114,55 @@ def http_quote(string):
             string = string.encode('ascii', 'replace')
     # Wrap in double-quotes for ; , and the like
     return '"{!s}"'.format(string.replace('\\', '\\\\').replace('"', '\\"'))
+
+
+try:
+    # From Django 1.4
+    from django.conf import override_settings
+except ImportError:
+    class override_settings(object):
+        """
+        Acts as either a decorator, or a context manager. If it's a decorator it
+        takes a function and returns a wrapped function. If it's a contextmanager
+        it's used with the ``with`` statement. In either event entering/exiting
+        are called before and after, respectively, the function/block is executed.
+        """
+        def __init__(self, **kwargs):
+            self.options = kwargs
+            self.wrapped = settings._wrapped
+
+        def __enter__(self):
+            self.enable()
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.disable()
+
+        def __call__(self, test_func):
+            from django.test import TransactionTestCase
+            if isinstance(test_func, type) and issubclass(test_func, TransactionTestCase):
+                original_pre_setup = test_func._pre_setup
+                original_post_teardown = test_func._post_teardown
+                def _pre_setup(innerself):
+                    self.enable()
+                    original_pre_setup(innerself)
+                def _post_teardown(innerself):
+                    original_post_teardown(innerself)
+                    self.disable()
+                test_func._pre_setup = _pre_setup
+                test_func._post_teardown = _post_teardown
+                return test_func
+            else:
+                @wraps(test_func)
+                def inner(*args, **kwargs):
+                    with self:
+                        return test_func(*args, **kwargs)
+            return inner
+
+        def enable(self):
+            override = copy(settings._wrapped)
+            for key, new_value in self.options.items():
+                setattr(override, key, new_value)
+            settings._wrapped = override
+
+        def disable(self):
+            settings._wrapped = self.wrapped
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):
author	Simon Law <simon.law@ecometrica.com>	2012-07-24 14:57:06 -0400
committer	Simon Law <simon.law@ecometrica.com>	2012-07-24 14:57:06 -0400
commit	8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7 (patch)
tree	e88c55a0ba32e5222547a2edb02a589a6cad6e84 /wkhtmltopdf
parent	4a27fe9f16ecd4bc6f94ad89046767450316490c (diff)
download	django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.tar.gz django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.tar.bz2 django-wkhtmltopdf-8e53c5bc5db462f6e39404c73ac96ec0cbb6a6c7.zip