Reliable command-line argument parsing for wkhtmltopdf().

The API for wkhtmltopdf has changed. Long arguments that take no
parameters now use True and not the empty string. In addition,
argument-parameters may now be Unicode.

2 files changed, 73 insertions, 20 deletions

diff --git a/wkhtmltopdf/tests.py b/wkhtmltopdf/tests.py
index f50d8f2..3e8359f 100644
--- a/wkhtmltopdf/tests.py
+++ b/wkhtmltopdf/tests.py
@@ -1,24 +1,58 @@
+# -*- coding: utf-8 -*-
+
 from __future__ import absolute_import
 
+from StringIO import StringIO
 import os
+import sys
 
 from django.test import TestCase
 
-from .utils import template_to_temp_file, wkhtmltopdf
+from .subprocess import CalledProcessError
+from .utils import _options_to_args, template_to_temp_file, wkhtmltopdf
 
 
 class TestUtils(TestCase):
+    def setUp(self):
+        # Clear standard error
+        self._stderr = sys.stderr
+        sys.stderr = open(os.devnull, 'w')
+
+    def tearDown(self):
+        sys.stderr = self._stderr
+
+    def test_options_to_args(self):
+        self.assertEqual(_options_to_args(), [])
+        self.assertEqual(_options_to_args(heart=u'♥', verbose=True,
+                                          file_name='file-name'),
+                         ['--file-name', 'file-name',
+                          '--heart', u'♥',
+                          '--verbose'])
+
     def test_wkhtmltopdf(self):
         """Should run wkhtmltopdf to generate a PDF"""
         title = 'A test template.'
         temp_file = template_to_temp_file('sample.html', {'title': title})
         pdf_output = None
         try:
+            # Standard call
             pdf_output = wkhtmltopdf(pages=[temp_file])
+            self.assertTrue(pdf_output.startswith('%PDF'), pdf_output)
+
+            # Single page
+            pdf_output = wkhtmltopdf(pages=temp_file)
+            self.assertTrue(pdf_output.startswith('%PDF'), pdf_output)
+
+            # Unicode
+            pdf_output = wkhtmltopdf(pages=[temp_file], title=u'♥')
+            self.assertTrue(pdf_output.startswith('%PDF'), pdf_output)
+
+            # Invalid arguments
+            self.assertRaises(CalledProcessError,
+                              wkhtmltopdf, pages=[])
         finally:
             if os.path.exists(temp_file):
                 os.remove(temp_file)
-        self.assertTrue(pdf_output.startswith('%PDF'), pdf_output)
 
     def test_template_to_temp_file(self):
         """Should render a template to a temporary file."""
diff --git a/wkhtmltopdf/utils.py b/wkhtmltopdf/utils.py
index 7994a39..aab4152 100644
--- a/wkhtmltopdf/utils.py
+++ b/wkhtmltopdf/utils.py
@@ -1,6 +1,8 @@
 from __future__ import absolute_import
 
+from itertools import chain
 from os import fdopen
+import sys
 from tempfile import mkstemp
 
 from django.conf import settings
@@ -11,12 +13,24 @@ from .subprocess import check_output
 
 WKHTMLTOPDF_CMD = getattr(settings, 'WKHTMLTOPDF_CMD', 'wkhtmltopdf')
 
+
+def _options_to_args(**options):
+    """Converts ``options`` into a string of command-line arguments."""
+    flags = []
+    for name in sorted(options):
+        value = options[name]
+        flags.append('--' + name.replace('_', '-'))
+        if value is not True:
+            flags.append(unicode(value))
+    return flags
+
+
 def wkhtmltopdf(pages, output=None, **kwargs):
     """
     Converts html to PDF using http://code.google.com/p/wkhtmltopdf/.
 
     pages: List of file paths or URLs of the html to be converted.
-    output: Optional output file path.
+    output: Optional output file path. If None, the output is returned.
     **kwargs: Passed to wkhtmltopdf via _extra_args() (See
               https://github.com/antialize/wkhtmltopdf/blob/master/README_WKHTMLTOPDF
               for acceptable args.)
@@ -24,31 +38,36 @@ def wkhtmltopdf(pages, output=None, **kwargs):
                   {'footer_html': 'http://example.com/foot.html'}
               becomes
                   '--footer-html http://example.com/foot.html'
-              Where there is no value passed, use a blank string. e.g.:
-                  {'disable_javascript': ''}
+              Where there is no value passed, use True. e.g.:
+                  {'disable_javascript': True}
               becomes:
-                  '--disable-javascript '
+                  '--disable-javascript'
 
     example usage:
-        wkhtmltopdf(html_path="~/example.html",
+        wkhtmltopdf(pages=['/tmp/example.html'],
                     dpi=300,
-                    orientation="Landscape",
-                    disable_javascript="")
+                    orientation='Landscape',
+                    disable_javascript=True)
     """
+    if isinstance(pages, basestring):
+        # Support a single page.
+        pages = [pages]
 
-    def _extra_args(**kwargs):
-        """Converts kwargs into a string of flags to be passed to wkhtmltopdf."""
-        flags = ''
-        for k, v in kwargs.items():
-            flags += ' --%s %s' % (k.replace('_', '-'), v)
-        return flags
+    if output is None:
+        # Standard output.
+        output = '-'
 
-    if not isinstance(pages, list):
-        pages = [pages]
+    # Default options:
+    options = {
+        'quiet': True,
+    }
+    options.update(kwargs)
 
-    kwargs['quiet'] = ''
-    args = '%s %s %s %s' % (WKHTMLTOPDF_CMD, _extra_args(**kwargs), ' '.join(pages), output or '-')
-    return check_output(args, shell=True)
+    args = list(chain([WKHTMLTOPDF_CMD],
+                      _options_to_args(**options),
+                      list(pages),
+                      [output]))
+    return check_output(args, stderr=sys.stderr)
 
 
 def template_to_temp_file(template_name, dictionary=None, context_instance=None):