diff options
| -rw-r--r-- | wkhtmltopdf/tests.py | 38 | ||||
| -rw-r--r-- | wkhtmltopdf/utils.py | 55 |
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): |