Divide PDF generation functionality in more functions to use separate parts

aleksis/core/util/pdf.py

@@ -2,7 +2,7 @@ import os
 import subprocess  # noqa
 from datetime import timedelta
 from tempfile import TemporaryDirectory
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Union, Callable
 from urllib.parse import urljoin
 
 from django.conf import settings
@@ -69,10 +69,8 @@ def generate_pdf(
     recorder.set_progress(1, 1)
 
 
-def generate_pdf_from_template(
-    template_name: str, context: Optional[dict] = None, request: Optional[HttpRequest] = None
-) -> Tuple[PDFFile, AsyncResult]:
-    """Start a PDF generation task and return the matching file object and Celery result."""
+def process_context_for_pdf(context: Optional[dict] = None, request: Optional[HttpRequest] = None):
+    context = context or {}
     if not request:
         processed_context = process_custom_context_processors(
             settings.NON_REQUEST_CONTEXT_PROCESSORS
@@ -80,10 +78,14 @@ def generate_pdf_from_template(
         processed_context.update(context)
     else:
         processed_context = context
-    html_template = render_to_string(template_name, processed_context, request)
+    return processed_context
+
+def generate_pdf_from_html(
+    html: str, request: Optional[HttpRequest] = None) -> Tuple[PDFFile, AsyncResult]:
+    """Start a PDF generation task and return the matching file object and Celery result."""
 
     file_object = PDFFile.objects.create(
-        html_file=ContentFile(html_template.encode(), name="source.html")
+        html_file=ContentFile(html.encode(), name="source.html")
     )
 
     # As this method may be run in background and there is no request available,
@@ -97,6 +99,19 @@ def generate_pdf_from_template(
 
     return file_object, result
 
+def generate_pdf_from_template(
+    template_name: str, context: Optional[dict] = None, request: Optional[HttpRequest] = None, render_method: Optional[Callable] = None
+) -> Tuple[PDFFile, AsyncResult]:
+    """Start a PDF generation task and return the matching file object and Celery result."""
+    processed_context = process_context_for_pdf(context, request)
+
+    if render_method:
+        html_template = render_method(processed_context, request)
+    else:
+        html_template = render_to_string(template_name, processed_context, request)
+
+    return generate_pdf_from_html(html_template, request)
+
 
 def render_pdf(
     request: Union[HttpRequest, None], template_name: str, context: dict = None