Add selenium to problem pdf (DMOJ)

dmoj/settings.py

@@ -125,6 +125,10 @@ SLIMERJS_PAPER_SIZE = 'Letter'
 PUPPETEER_MODULE = '/usr/lib/node_modules/puppeteer'
 PUPPETEER_PAPER_SIZE = 'Letter'
 
+USE_SELENIUM = False
+SELENIUM_CUSTOM_CHROME_PATH = None
+SELENIUM_CHROMEDRIVER_PATH = 'chromedriver'
+
 PYGMENT_THEME = 'pygment-github.css'
 INLINE_JQUERY = True
 INLINE_FONTAWESOME = True

judge/management/commands/render_pdf.py

@@ -8,8 +8,8 @@ from django.template.loader import get_template
 from django.utils import translation
 
 from judge.models import Problem, ProblemTranslation
-from judge.pdf_problems import DefaultPdfMaker, PhantomJSPdfMaker, PuppeteerPDFRender, SlimerJSPdfMaker
-
+from judge.pdf_problems import DefaultPdfMaker, PhantomJSPdfMaker, PuppeteerPDFRender, SeleniumPDFRender, \
+    SlimerJSPdfMaker
 
 class Command(BaseCommand):
     help = 'renders a PDF file of a problem'
@@ -24,6 +24,7 @@ class Command(BaseCommand):
         parser.add_argument('-s', '--slimerjs', action='store_const', const=SlimerJSPdfMaker, dest='engine')
         parser.add_argument('-c', '--chrome', '--puppeteer', action='store_const',
                             const=PuppeteerPDFRender, dest='engine')
+        parser.add_argument('-S', '--selenium', action='store_const', const=SeleniumPDFRender, dest='engine')
 
     def handle(self, *args, **options):
         try:

judge/pdf_problems.py

@@ -1,3 +1,4 @@
+import base64
 import errno
 import io
 import json
@@ -10,6 +11,20 @@ import uuid
 from django.conf import settings
 from django.utils.translation import gettext
 
+logger = logging.getLogger('judge.problem.pdf')
+
+HAS_SELENIUM = False
+if settings.USE_SELENIUM:
+    try:
+        from selenium import webdriver
+        from selenium.common.exceptions import TimeoutException
+        from selenium.webdriver.common.by import By
+        from selenium.webdriver.support import expected_conditions as EC
+        from selenium.webdriver.support.ui import WebDriverWait
+        HAS_SELENIUM = True
+    except ImportError:
+        logger.warning('Failed to import Selenium', exc_info=True)
+
 HAS_PHANTOMJS = os.access(settings.PHANTOMJS, os.X_OK)
 HAS_SLIMERJS = os.access(settings.SLIMERJS, os.X_OK)
 
@@ -18,13 +33,11 @@ PUPPETEER_MODULE = settings.PUPPETEER_MODULE
 HAS_PUPPETEER = os.access(NODE_PATH, os.X_OK) and os.path.isdir(PUPPETEER_MODULE)
 
 HAS_PDF = (os.path.isdir(settings.DMOJ_PDF_PROBLEM_CACHE) and
-           (HAS_PHANTOMJS or HAS_SLIMERJS or HAS_PUPPETEER))
+           (HAS_PHANTOMJS or HAS_SLIMERJS or HAS_PUPPETEER or HAS_SELENIUM))
 
 EXIFTOOL = settings.EXIFTOOL
 HAS_EXIFTOOL = os.access(EXIFTOOL, os.X_OK)
 
-logger = logging.getLogger('judge.problem.pdf')
-
 
 class BasePdfMaker(object):
     math_engine = 'jax'
@@ -240,8 +253,8 @@ puppeteer.launch().then(browser => Promise.resolve()
 
     def get_render_script(self):
         return self.template.replace('{params}', json.dumps({
-            'input': 'file://' + os.path.abspath(os.path.join(self.dir, 'input.html')),
-            'output': os.path.abspath(os.path.join(self.dir, 'output.pdf')),
+            'input': 'file://%s' % self.htmlfile,
+            'output': self.pdffile,
             'paper': settings.PUPPETEER_PAPER_SIZE,
             'footer': gettext('Page [page] of [topage]'),
         }))
@@ -257,9 +270,51 @@ puppeteer.launch().then(browser => Promise.resolve()
         self.proc = subprocess.Popen(cmdline, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=self.dir, env=env)
         self.log = self.proc.communicate()[0]
 
+class SeleniumPDFRender(BasePdfMaker):
+    success = False
+    template = {
+        'printBackground': True,
+        'displayHeaderFooter': True,
+        'headerTemplate': '<div></div>',
+        'footerTemplate': '<center style="margin: 0 auto; font-family: Segoe UI; font-size: 10px">' +
+                          gettext('Page %s of %s') %
+                          ('<span class="pageNumber"></span>', '<span class="totalPages"></span>') +
+                          '</center>',
+    }
+
+    def get_log(self, driver):
+        return '\n'.join(map(str, driver.get_log('driver') + driver.get_log('browser')))
+
+    def _make(self, debug):
+        options = webdriver.ChromeOptions()
+        options.add_argument("--headless")
+        options.binary_location = settings.SELENIUM_CUSTOM_CHROME_PATH
+
+        browser = webdriver.Chrome(settings.SELENIUM_CHROMEDRIVER_PATH, options=options)
+        browser.get('file://%s' % self.htmlfile)
+        self.log = self.get_log(browser)
+
+        try:
+            WebDriverWait(browser, 15).until(EC.presence_of_element_located((By.CLASS_NAME, 'math-loaded')))
+        except TimeoutException:
+            logger.error('PDF math rendering timed out')
+            self.log = self.get_log(browser) + '\nPDF math rendering timed out'
+            return
+        response = browser.execute_cdp_cmd('Page.printToPDF', self.template)
+        self.log = self.get_log(browser)
+        if not response:
+            return
+
+        with open(self.pdffile, 'wb') as f:
+            f.write(base64.b64decode(response['data']))
+
+        self.success = True
+
 
 if HAS_PUPPETEER:
     DefaultPdfMaker = PuppeteerPDFRender
+elif HAS_SELENIUM:
+    DefaultPdfMaker = SeleniumPDFRender
 elif HAS_SLIMERJS:
     DefaultPdfMaker = SlimerJSPdfMaker
 elif HAS_PHANTOMJS:

judge/views/problem.py

@@ -256,7 +256,6 @@ class ProblemPdfView(ProblemMixin, SingleObjectMixin, View):
                     'math_engine': maker.math_engine,
                 }).replace('"//', '"https://').replace("'//", "'https://")
                 maker.title = problem_name
-
                 assets = ['style.css', 'pygment-github.css']
                 if maker.math_engine == 'jax':
                     assets.append('mathjax_config.js')
@@ -267,7 +266,6 @@ class ProblemPdfView(ProblemMixin, SingleObjectMixin, View):
                     self.logger.error('Failed to render PDF for %s', problem.code)
                     return HttpResponse(maker.log, status=500, content_type='text/plain')
                 shutil.move(maker.pdffile, cache)
-
         response = HttpResponse()
         if hasattr(settings, 'DMOJ_PDF_PROBLEM_INTERNAL') and \
                 request.META.get('SERVER_SOFTWARE', '').startswith('nginx/'):