chrome_pdf_generator.py

"""
Courtesy of https://github.com/loftylabs/django-hardcopy
"""
import platform
import subprocess
from pathlib import Path
from tempfile import NamedTemporaryFile

from reseller import app

LINUX_PATHS = [
    '/usr/bin/chromium',
    '/usr/bin/chromium-browser',
    '/usr/bin/chrome',
    '/usr/bin/google-chrome',
    '/usr/bin/chrome-browser',
]


class PdfGenerator:
    chrome_path = getattr(app.config, 'CHROME_PATH', None)

    chrome_args = []
    chrome_kwargs = {}

    def __init__(self, html, pdf_file=None, root_dir=None):
        self.input_file = NamedTemporaryFile(suffix='.html')
        if pdf_file is None:
            self.output_file = NamedTemporaryFile(dir=root_dir)
        else:
            self.output_file = pdf_file

        if isinstance(html, str):
            html = bytes(html.encode('utf-8'))
        self.input_file.write(html)
        self.input_file.flush()

    def get_chrome_path(self):
        if not self.chrome_path:
            self.chrome_path = PdfGenerator.guess_chrome_path()
        return self.chrome_path

    def get_chrome_args(self):
        default_args = [
            self.get_chrome_path(),
            '--no-sandbox',  # Avoids permission issues while dockerized.
            '--headless',
            '--disable-extensions',  # Reduces startup overhead.
            '--disable-gpu',  # Required by chrome's headless mode for now.
            f'--print-to-pdf="{self.output_file.name}"',
            f'file://{self.input_file.name}',
        ]
        args_without_values = [f'--{arg}' for arg in self.chrome_args]
        args_with_values = [f'--{k}={v}' for k, v in self.chrome_kwargs.items()]

        default_args += args_without_values
        default_args += args_with_values

        return default_args

    def call_chrome(self):
        chrome_args = self.get_chrome_args()
        subprocess.call(" ".join(chrome_args), shell=True)
        self.output_file.seek(0)

    def generate_pdf(self):
        self.call_chrome()
        return self.output_file.read()

    @staticmethod
    def guess_chrome_path():
        """Attempt to guess the Chrome path by default."""
        if platform.uname()[0] == "Darwin":
            # pylint: disable=anomalous-backslash-in-string
            return '/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome'
        if platform.uname()[0] == "Linux":
            # Iterate through some sane path defaults.
            for path in LINUX_PATHS:
                if Path(path).is_file():
                    return path
        # No path found, throw an error.
        raise ValueError('Missing CHROME_PATH! Unable to resolve path!')