blog.bythewood.me-wagtail

  1"""
  2A simple library for working with the chromium browser in headless mode for
  3generating things like screenshots.
  4"""
  5import os
  6import shutil
  7import subprocess
  8import uuid
  9
 10from django.core.files.storage import default_storage
 11
 12
 13# Get chromium path, it's sometimes chromium and sometimes chromium-browser
 14chromium = None
 15if shutil.which("chromium"):
 16    chromium = "chromium"
 17elif shutil.which("chromium-browser"):
 18    chromium = "chromium-browser"
 19
 20
 21base_command = [
 22    chromium,
 23    "--headless",
 24    "--no-sandbox",
 25    "--use-gl=swiftshader",
 26    "--disable-gpu",
 27    "--disable-software-rasterizer",
 28    "--disable-dev-shm-usage",
 29    "--disable-crash-reporter",
 30    "--disable-extensions",
 31    "--disable-in-process-stack-traces",
 32    "--disable-logging",
 33    "--window-size=1280x720",
 34    "--hide-scrollbars",
 35]
 36
 37
 38def save_tempfile_to_storage(tempfilename, filename):
 39    """
 40    Saves the given tempfile to django default_storage.
 41
 42    :param tempfilename: The tempfile we want to save
 43    :param filename: The storage location to save the file to
 44    """
 45    if default_storage.exists(filename):
 46        default_storage.delete(filename)
 47    default_storage.save(filename, open(tempfilename, "rb"))
 48    os.remove(tempfilename)
 49
 50
 51def run_chromium_command(command):
 52    """
 53    Runs the given chromium command and returns the stdout.
 54
 55    :param command: The command to run
 56    """
 57    command = command.split()
 58    command = base_command + command
 59    subprocess.run(
 60        command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
 61    )
 62
 63
 64def generate_screenshot_from_url(url, filename):
 65    """
 66    Generates a screenshot of the given url and saves it to the given output
 67    file.
 68
 69    :param url: The url to screenshot
 70    :param filename: The output file to save the screenshot to
 71    """
 72    tempfilename = f"{uuid.uuid4()}.png"
 73    run_chromium_command(f"--screenshot={tempfilename} {url}")
 74    save_tempfile_to_storage(tempfilename, filename)
 75    return default_storage.url(filename)
 76
 77
 78def generate_screenshot_from_html(html, filename):
 79    """
 80    Generates a screenshot of the given html and saves it to the given output
 81    file.
 82
 83    :param html: The html to screenshot
 84    :param filename: The output file to save the screenshot to
 85    """
 86    tempfilename = f"{uuid.uuid4()}.html"
 87    with open(tempfilename, "w") as f:
 88        f.write(html)
 89    tempfilename_path = "file://" + os.path.join(os.getcwd(), tempfilename)
 90    run_chromium_command(f"--screenshot={tempfilename} {tempfilename_path}")
 91    save_tempfile_to_storage(tempfilename, filename)
 92    return default_storage.url(filename)
 93
 94
 95def generate_pdf_from_url(url, filename):
 96    """
 97    Generates a pdf of the given url and saves it to the given output file.
 98
 99    :param url: The url to screenshot
100    :param filename: The output file to save the screenshot to
101    """
102    tempfilename = f"{uuid.uuid4()}.pdf"
103    run_chromium_command(f"--print-to-pdf-no-header --print-to-pdf={tempfilename} {url}")
104    save_tempfile_to_storage(tempfilename, filename)
105    return default_storage.url(filename)
106
107
108def generate_pdf_from_html(html, filename):
109    """
110    Generates a pdf of the given html and saves it to the given output file.
111
112    :param url: The url to screenshot
113    :param filename: The output file to save the screenshot to
114    """
115    tempfilename = f"{uuid.uuid4()}.html"
116    with open(tempfilename, "w") as f:
117        f.write(html)
118    tempfilename_path = "file://" + os.path.join(os.getcwd(), tempfilename)
119    run_chromium_command(f"--print-to-pdf-no-header --print-to-pdf={tempfilename} {tempfilename_path}")
120    save_tempfile_to_storage(tempfilename, filename)
121    return default_storage.url(filename)