Merge pull request #104 from halfnibble/master

Refactor to better support multiple pages.

3 files changed, 100 insertions, 40 deletions

diff --git a/wkhtmltopdf/tests/run.py b/wkhtmltopdf/tests/run.py
index 5004ee1..e0021ab 100644
--- a/wkhtmltopdf/tests/run.py
+++ b/wkhtmltopdf/tests/run.py
@@ -10,7 +10,7 @@ DIRNAME = os.path.abspath(os.path.dirname(__file__))
 sys.path.insert(0, os.getcwd())
 
 settings.configure(
-    DEBUG=True,
+    DEBUG=False,
     DATABASES={
         'default': {
             'ENGINE': 'django.db.backends.sqlite3',
@@ -29,7 +29,7 @@ settings.configure(
     MEDIA_URL='/media/',
     STATIC_ROOT=os.path.join(DIRNAME, 'static'),
     STATIC_URL='/static/',
-    WKHTMLTOPDF_DEBUG=True,
+    WKHTMLTOPDF_DEBUG=False,
 )
 
 try:
diff --git a/wkhtmltopdf/tests/tests.py b/wkhtmltopdf/tests/tests.py
index 822f529..7b48421 100644
--- a/wkhtmltopdf/tests/tests.py
+++ b/wkhtmltopdf/tests/tests.py
@@ -14,7 +14,8 @@ from django.utils.encoding import smart_str
 
 from wkhtmltopdf.subprocess import CalledProcessError
 from wkhtmltopdf.utils import (_options_to_args, make_absolute_paths,
-                               wkhtmltopdf, render_to_temporary_file)
+                               wkhtmltopdf, render_to_temporary_file,
+                               RenderedFile)
 from wkhtmltopdf.views import PDFResponse, PDFTemplateView, PDFTemplateResponse
 
 
@@ -30,6 +31,7 @@ class UnicodeContentPDFTemplateView(PDFTemplateView):
         context['title'] = u'♥'
         return context
 
+
 class TestUtils(TestCase):
     def setUp(self):
         # Clear standard error
@@ -94,6 +96,45 @@ class TestUtils(TestCase):
         self.assertTrue(title in saved_content)
         temp_file.close()
 
+    def _render_file(self, template, context):
+        """Helper method for testing rendered file deleted/persists tests."""
+        render = RenderedFile(template=template, context=context)
+        render.temporary_file.seek(0)
+        saved_content = smart_str(render.temporary_file.read())
+
+        return (saved_content, render.filename)
+
+    def test_rendered_file_deleted_on_production(self):
+        """If WKHTMLTOPDF_DEBUG=False, delete rendered file on object close."""
+        title = 'A test template.'
+        template = loader.get_template('sample.html')
+        debug = getattr(settings, 'WKHTMLTOPDF_DEBUG', settings.DEBUG)
+
+        saved_content, filename = self._render_file(template=template,
+                                                    context={'title': title})
+        # First verify temp file was rendered correctly.
+        self.assertTrue(title in saved_content)
+
+        # Then check if file is deleted when debug=False.
+        self.assertFalse(debug)
+        self.assertFalse(os.path.isfile(filename))
+
+    def test_rendered_file_persists_on_debug(self):
+        """If WKHTMLTOPDF_DEBUG=True, the rendered file should persist."""
+        title = 'A test template.'
+        template = loader.get_template('sample.html')
+        with self.settings(WKHTMLTOPDF_DEBUG=True):
+            debug = getattr(settings, 'WKHTMLTOPDF_DEBUG', settings.DEBUG)
+
+            saved_content, filename = self._render_file(template=template,
+                                                    context={'title': title})
+            # First verify temp file was rendered correctly.
+            self.assertTrue(title in saved_content)
+
+            # Then check if file persists when debug=True.
+            self.assertTrue(debug)
+            self.assertTrue(os.path.isfile(filename))
+
 
 class TestViews(TestCase):
     template = 'sample.html'
diff --git a/wkhtmltopdf/utils.py b/wkhtmltopdf/utils.py
index 1866b4a..c871bd0 100644
--- a/wkhtmltopdf/utils.py
+++ b/wkhtmltopdf/utils.py
@@ -113,58 +113,77 @@ def convert_to_pdf(filename, header_filename=None, footer_filename=None, cmd_opt
     # Clobber header_html and footer_html only if filenames are
     # provided. These keys may be in self.cmd_options as hardcoded
     # static files.
+    # The argument `filename` may be a string or a list. However, wkhtmltopdf
+    # will coerce it into a list if a string is passed.
     cmd_options = cmd_options if cmd_options else {}
 
     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)
+    return wkhtmltopdf(pages=filename, **cmd_options)
 
-def render_pdf_from_template(input_template, header_template, footer_template, context, request=None, cmd_options=None):
-    debug = getattr(settings, 'WKHTMLTOPDF_DEBUG', settings.DEBUG)
-    cmd_options = cmd_options if cmd_options else {}
+class RenderedFile(object):
+    """
+    Create a temporary file resource of the rendered template with context.
+    The filename will be used for later conversion to PDF.
+    """
+    temporary_file = None
+    filename = ''
 
-    input_file = header_file = footer_file = None
-    header_filename = footer_filename = None
+    def __init__(self, template, context, request=None):
+        debug = getattr(settings, 'WKHTMLTOPDF_DEBUG', settings.DEBUG)
 
-    try:
-        input_file = render_to_temporary_file(
-            template=input_template,
+        self.temporary_file = render_to_temporary_file(
+            template=template,
             context=context,
             request=request,
             prefix='wkhtmltopdf', suffix='.html',
             delete=(not debug)
         )
+        self.filename = self.temporary_file.name
+
+    def __del__(self):
+        # Always close the temporary_file on object destruction.
+        if self.temporary_file is not None:
+            self.temporary_file.close()
+
+def render_pdf_from_template(input_template, header_template, footer_template, context, request=None, cmd_options=None):
+    # For basic usage. Performs all the actions necessary to create a single
+    # page PDF from a single template and context.
+    cmd_options = cmd_options if cmd_options else {}
+
+    header_filename = footer_filename = None
+
+    # Main content.
+    input_file = RenderedFile(
+        template=input_template,
+        context=context,
+        request=request
+    )
+
+    # Optional. For header template argument.
+    if header_template:
+        header_file = RenderedFile(
+            template=header_template,
+            context=context,
+            request=request
+        )
+        header_filename = header_file.filename
+
+    # Optional. For footer template argument.
+    if footer_template:
+        footer_file = RenderedFile(
+            template=footer_template,
+            context=context,
+            request=request
+        )
+        footer_filename = footer_file.filename
 
-        if header_template:
-            header_file = render_to_temporary_file(
-                template=header_template,
-                context=context,
-                request=request,
-                prefix='wkhtmltopdf', suffix='.html',
-                delete=(not debug)
-            )
-            header_filename = header_file.name
-
-        if footer_template:
-            footer_file = render_to_temporary_file(
-                template=footer_template,
-                context=context,
-                request=request,
-                prefix='wkhtmltopdf', suffix='.html',
-                delete=(not debug)
-            )
-            footer_filename = footer_file.name
-
-        return convert_to_pdf(filename=input_file.name,
-                              header_filename=header_filename,
-                              footer_filename=footer_filename,
-                              cmd_options=cmd_options)
-    finally:
-        # Clean up temporary files
-        for f in filter(None, (input_file, header_file, footer_file)):
-            f.close()
+    return convert_to_pdf(filename=input_file.filename,
+                          header_filename=header_filename,
+                          footer_filename=footer_filename,
+                          cmd_options=cmd_options)
 
 def content_disposition_filename(filename):
     """