From a6cf3e24fc76cb4471d6e0bfae140af8f98a65bb Mon Sep 17 00:00:00 2001
From: Shuhei Iitsuka <tushuhei@gmail.com>
Date: Thu, 5 Dec 2024 06:44:39 +0900
Subject: [PATCH] Add PDF option (#22)

---
 sphinxcontrib/screenshot.py | 26 ++++++++++++++++++++++++--
 tests/test-root/index.rst   |  5 +++++
 tests/test_it.py            | 17 +++++++++++++++--
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/sphinxcontrib/screenshot.py b/sphinxcontrib/screenshot.py
index ce181d3..414625c 100644
--- a/sphinxcontrib/screenshot.py
+++ b/sphinxcontrib/screenshot.py
@@ -72,6 +72,21 @@ class ScreenshotDirective(SphinxDirective):
 
     document.querySelector('button').click();
   ```
+
+  Use `figclass` option if you want to specify a class name to the image.
+
+  ```rst
+  .. screenshot:: http://www.example.com
+    :figclass: foo
+  ```
+
+  It also generates a PDF file when `pdf` option is given, which might be
+  useful when you need scalable image assets.
+
+  ```rst
+  .. screenshot:: http://www.example.com
+    :pdf:
+  ```
   """
 
   required_arguments = 1  # URL
@@ -81,12 +96,13 @@ class ScreenshotDirective(SphinxDirective):
       'width': directives.positive_int,
       'caption': directives.unchanged,
       'figclass': directives.unchanged,
+      'pdf': directives.flag,
   }
   pool = ThreadPoolExecutor()
 
   @staticmethod
   def take_screenshot(url: str, width: int, height: int, filepath: str,
-                      init_script: str, interactions: str):
+                      init_script: str, interactions: str, generate_pdf: bool):
     """Takes a screenshot with Playwright's Chromium browser.
 
     Args:
@@ -99,6 +115,7 @@ def take_screenshot(url: str, width: int, height: int, filepath: str,
         https://playwright.dev/python/docs/api/class-page#page-add-init-script
       interactions (str): JavaScript code to run before taking the screenshot
         after the page was loaded.
+      generate_pdf (bool): Generate a PDF file along with the screenshot.
     """
     with sync_playwright() as playwright:
       browser = playwright.chromium.launch()
@@ -119,6 +136,10 @@ def take_screenshot(url: str, width: int, height: int, filepath: str,
         raise RuntimeError('Timeout error occured at %s in executing\n%s' %
                            (url, interactions))
       page.screenshot(path=filepath)
+      if generate_pdf:
+        page.emulate_media(media='screen')
+        root, ext = os.path.splitext(filepath)
+        page.pdf(width=f'{width}px', height=f'{height}px', path=root + '.pdf')
       page.close()
       browser.close()
 
@@ -135,6 +156,7 @@ def run(self) -> typing.List[nodes.Node]:
     width = self.options.get('width', 1280)
     caption_text = self.options.get('caption', '')
     figclass = self.options.get('figclass', '')
+    pdf = 'pdf' in self.options
     interactions = '\n'.join(self.content)
 
     if urlparse(url).scheme not in {'http', 'https'}:
@@ -150,7 +172,7 @@ def run(self) -> typing.List[nodes.Node]:
     if not os.path.exists(filepath):
       fut = self.pool.submit(ScreenshotDirective.take_screenshot, url, width,
                              height, filepath, screenshot_init_script,
-                             interactions)
+                             interactions, pdf)
       fut.result()
 
     # Create image and figure nodes
diff --git a/tests/test-root/index.rst b/tests/test-root/index.rst
index af6f02f..82992e5 100644
--- a/tests/test-root/index.rst
+++ b/tests/test-root/index.rst
@@ -19,3 +19,8 @@
   :width: 480
   :height: 320
   :figclass: round
+
+.. screenshot:: http://www.example.com
+  :width: 640
+  :height: 640
+  :pdf:
diff --git a/tests/test_it.py b/tests/test_it.py
index d9fb4f7..835ea13 100644
--- a/tests/test_it.py
+++ b/tests/test_it.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from io import StringIO
 
 import pytest
@@ -29,7 +30,7 @@ def test_default(app: SphinxTestApp, status: StringIO,
 
   # Every screenshot directive should become an image.
   imgs = soup.find_all('img')
-  assert len(list(imgs)) == 4
+  assert len(list(imgs)) == 5
 
   # The image size should be set as specified.
   img_obj = Image.open(app.outdir / imgs[0]['src'])
@@ -54,4 +55,16 @@ def test_default(app: SphinxTestApp, status: StringIO,
   assert img_with_caption_a['src'] == img_with_caption_b['src']
 
   # The figure node should have the class name specified.
-  assert 'round' in soup.find_all('figure')[-1]['class']
+  assert 'round' in soup.find_all('figure')[3]['class']
+
+  # Should generate a PDF file if specified.
+  img_with_pdf = imgs[4]['src']
+  root, ext = os.path.splitext(os.path.basename(img_with_pdf))
+  pdf_filepath = app.outdir / '_static' / 'screenshots' / f'{root}.pdf'
+  assert os.path.exists(pdf_filepath)
+
+  # Should not generate a PDF file if not specified.
+  img_without_pdf = imgs[2]['src']
+  root, ext = os.path.splitext(os.path.basename(img_without_pdf))
+  pdf_filepath = app.outdir / '_static' / 'screenshots' / f'{root}.pdf'
+  assert not os.path.exists(pdf_filepath)