Refactor the main render method into a class based view

Use a PdfResponse to deal with the headers and response type.
Make the margin_* and filename variables instance variables on the view
with sane defaults.

2 files changed, 62 insertions, 30 deletions

diff --git a/wkhtmltopdf/utils.py b/wkhtmltopdf/utils.py
index b098383..4666f4c 100644
--- a/wkhtmltopdf/utils.py
+++ b/wkhtmltopdf/utils.py
@@ -1,9 +1,8 @@
-from os import fdopen, remove
+from os import fdopen
 from subprocess import Popen, PIPE, CalledProcessError
 from tempfile import mkstemp
 
 from django.conf import settings
-from django.http import HttpResponse
 from django.template import loader
 from django.utils.encoding import smart_str
 
@@ -58,33 +57,6 @@ def wkhtmltopdf(pages, output=None, **kwargs):
 
     return output
 
-def render_to_pdf(template_name, dictionary=None, context_instance=None,
-                    header_template=None, footer_template=None, **kwargs):
-    """Renders a html template as a PDF response."""
-
-    filename = kwargs.pop('filename', None)
-
-    page_path = template_to_temp_file(template_name, dictionary, context_instance)
-
-    tmp_files = []
-    if header_template is not None:
-        kwargs['header_html'] = template_to_temp_file(header_template, dictionary, context_instance)
-        tmp_files.append(kwargs['header_html'])
-    if footer_template is not None:
-        kwargs['footer_html'] = template_to_temp_file(footer_template, dictionary, context_instance)
-        tmp_files.append(kwargs['footer_html'])
-
-    content = wkhtmltopdf([page_path], **kwargs)
-
-    for path in tmp_files:
-        remove(path)
-
-    response = HttpResponse(content, mimetype='application/pdf')
-    if filename is not None:
-        response['Content-Disposition'] = 'attachment;' + ' filename=%s' %filename
-    return response
-
-
 def template_to_temp_file(*args, **kwargs):
     """
     Renders a template to a temp file, and returns the path of the file.
@@ -94,4 +66,3 @@ def template_to_temp_file(*args, **kwargs):
         f.write(smart_str(loader.render_to_string(*args, **kwargs)))
     return tempfile_path
 
-
diff --git a/wkhtmltopdf/views.py b/wkhtmltopdf/views.py
new file mode 100644
index 0000000..e2fee19
--- /dev/null
+++ b/wkhtmltopdf/views.py
@@ -0,0 +1,61 @@
+from os import remove
+from re import compile
+
+from django.conf import settings
+from django.contrib.sites.models import Site
+from django.template.response import HttpResponse
+from django.views.generic import TemplateView
+
+from wkhtmltopdf.utils import template_to_temp_file, wkhtmltopdf
+
+class PdfResponse(HttpResponse):
+    def __init__(self, content, filename):
+        super(PdfResponse, self).__init__(content, 'application/pdf')
+        self.__setitem__('Content-Disposition', 'attachment; filename=%s' % filename)
+
+
+class PdfTemplateView(TemplateView):
+    filename = 'rendered_pdf.pdf'
+    footer_template = None
+    header_template = None
+    margin_bottom = 0
+    margin_left = 0
+    margin_right = 0
+    margin_top = 0
+    response = PdfResponse
+
+    def get(self, request, context_instance=None, *args, **kwargs):
+        if request.GET.get('as', '') == 'html':
+            super(PdfTemplateView, self).get(request, *args, **kwargs)
+
+        page_path = template_to_temp_file(self.template_name, self.get_context_data(), context_instance)
+
+        tmp_files = []
+        if self.header_template:
+            kwargs['header_html'] = template_to_temp_file(self.header_template, self.get_context_data(), context_instance)
+            tmp_files.append(kwargs['header_html'])
+        if self.footer_template:
+            kwargs['footer_html'] = template_to_temp_file(self.footer_template, self.get_context_data(), context_instance)
+            tmp_files.append(kwargs['footer_html'])
+
+        map(remove, tmp_files)
+
+        kwargs.update({
+            'margin_bottom': self.margin_bottom,
+            'margin_left': self.margin_left,
+            'margin_right': self.margin_right,
+            'margin_top': self.margin_top
+        })
+        return self.response(wkhtmltopdf(page_path, **kwargs), self.filename)
+
+    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.STATIC_URL):
+            context['MEDIA_URL'] = 'http://' + Site.objects.get_current().domain + settings.MEDIA_URL
+
+        return context
+