feat: the 'screenshot' directive inherits from 'figure'

tushuhei · Jan 4, 2025 · a56fa02 · a56fa02
1 parent 42d452d
commit a56fa02
Show file tree

Hide file tree

Showing 10 changed files with 101 additions and 168 deletions.
diff --git a/README.md b/README.md
@@ -5,8 +5,8 @@ A Sphinx extension to embed website screenshots.
 ```rst
 .. screenshot:: http://www.example.com
   :browser: chromium
-  :width: 1280
-  :height: 960
+  :viewport-width: 1280
+  :viewport-height: 960
   :color-scheme: dark
 ```
 

diff --git a/doc/index.rst b/doc/index.rst
@@ -2,9 +2,8 @@
 sphinxcontrib-screenshot
 ========================
 
-A Sphinx extension to embed website screenshots.
 
-.. screenshot:: https://github.com/tushuhei/sphinxcontrib-screenshot
+A Sphinx extension to embed website screenshots.
 
 Install
 #######
@@ -35,13 +34,22 @@ Then use the `screenshot` directive in your Sphinx source file.
 Options
 #######
 
-You can describe the interaction that you want to have with the webpage before taking a screenshot in JavaScript.
+.. screenshot:: https://github.com/tushuhei/sphinxcontrib-screenshot
+    :align: right
+    :width: 400
+
+    An example of screenshot using the figure `:align:` and `:width:` options.
+
+`screenshot` inherits from the `figure directive <https://docutils.sourceforge.io/docs/ref/rst/directives.html#image>`__
+and supports all its options (`:align:`, `:alt:`, `:figclass:`, `:figwidth:`, `:height:`, `:loading:`, `:scale:`, `:target:`, `:width:`)
 
 .. code-block:: rst
 
-    .. screenshot:: http://www.example.com
+    .. screenshot:: https://github.com/tushuhei/sphinxcontrib-screenshot
+        :align: right
+        :width: 400
 
-      document.querySelector('button').click();
+        An example of screenshot using the figure `:align:` and `:width:` options.
 
 .. _browser:
 
@@ -55,21 +63,6 @@ You can choose the browser to use to take the screenshots with the :code:`:brows
     .. screenshot:: http://www.example.com
       :browser: firefox
 
-.. _caption:
-
-``:caption:``
-=============
-
-You can include a caption for the screenshot's :code:`figure` directive.
-
-.. code-block:: rst
-
-   .. screenshot:: http://www.example.com
-    :caption: This is a screenshot for www.example.com
-
-.. screenshot:: http://www.example.com
-  :caption: This is a screenshot for www.example.com
-
 .. _color-scheme:
 
 ``:color-scheme:``
@@ -94,18 +87,6 @@ The custom context to use for taking the screenshot. See :ref:`screenshot_contex
     .. screenshot:: http://www.example.com
       :context: logged-as-user
 
-.. _figclass:
-
-``:figclass:``
-==============
-
-Use :code:`figclass` option if you want to specify a class name to the image.
-
-.. code-block:: rst
-
-    .. screenshot:: http://www.example.com
-      :figclass: foo
-
 .. _full-page:
 
 ``:full-page:``
@@ -132,6 +113,19 @@ You can pass additional headers to the requests, for instance to customize the d
         Authorization Bearer my-super-secret-token
         Accept-Language fr-FR,fr
 
+.. _interactions:
+
+``:interactions:``
+==================
+
+You can describe the interaction that you want to have with the webpage before taking a screenshot in JavaScript.
+
+.. code-block:: rst
+
+    .. screenshot:: http://www.example.com
+      :interactions:
+        document.querySelector('button').click();
+
 .. _pdf:
 
 ``:pdf:``
@@ -151,20 +145,20 @@ It also generates a PDF file when :code:`pdf` option is given, which might be us
 .. _width:
 .. _height:
 
-``:width:`` and ``:height:``
-============================
+``:viewport-width:`` and ``:viewport-height:``
+==============================================
 
 You can specify the screen size for a particular screenshot with :code:`width` and :code:`height` parameters.
 
 .. code-block:: rst
 
     .. screenshot:: http://www.example.com
-      :width: 800
-      :height: 600
+      :viewport-width: 800
+      :viewport-height: 600
 
 .. screenshot:: http://www.example.com
-  :width: 800
-  :height: 600
+  :viewport-width: 800
+  :viewport-height: 600
 
 Configuration
 #############
@@ -257,18 +251,18 @@ Those are the default headers to be used when taking screenshots. They can be ov
         "Accept-Language": "fr-FR,fr",
     }
 
-.. _screenshot_default_width:
-.. _screenshot_default_height:
+.. _screenshot_default_viewport_width:
+.. _screenshot_default_viewport_height:
 
-``screenshot_default_width`` and ``screenshot_default_height``
-==============================================================
+``screenshot_default_viewport_width`` and ``screenshot_default_viewport_height``
+================================================================================
 
-You can define the default size of your screenshots in `conf.py`, those values will be used by default when :ref:`:width: <width>` and :ref:`:height: <height>` are not set:
+You can define the default size of your screenshots in `conf.py`, those values will be used by default when :ref:`:viewport-width: <width>` and :ref:`:viewport-height: <height>` are not set:
 
 .. code-block:: python
 
-    screenshot_default_width = 1920
-    screenshot_default_height = 1200
+    screenshot_default_viewport_width = 1920
+    screenshot_default_viewport_height = 1200
 
 Local WSGI application
 ######################

diff --git a/sphinxcontrib/screenshot.py b/sphinxcontrib/screenshot.py
@@ -24,7 +24,7 @@
 
 from docutils import nodes
 from docutils.parsers.rst import directives
-from docutils.statemachine import ViewList
+from docutils.parsers.rst.directives.images import Figure
 from playwright._impl._helper import ColorScheme
 from playwright.sync_api import Browser, BrowserContext
 from playwright.sync_api import TimeoutError as PlaywrightTimeoutError
@@ -41,7 +41,7 @@
 })
 
 
-class ScreenshotDirective(SphinxDirective):
+class ScreenshotDirective(SphinxDirective, Figure):
   """Sphinx Screenshot Dirctive.
 
   This directive embeds a screenshot of a webpage.
@@ -54,20 +54,13 @@ class ScreenshotDirective(SphinxDirective):
   .. screenshot:: http://www.example.com
   ```
 
-  You can also specify the screen size for the screenshot with `width` and
-  `height` parameters in pixel.
+  You can also specify the screen size for the screenshot with
+  `viewport-width` and `viewport-height` parameters in pixel.
 
   ```rst
   .. screenshot:: http://www.example.com
-    :width: 1280
-    :height: 960
-  ```
-
-  You can include a caption for the screenshot's `figure` directive.
-
-  ```rst
-  .. screenshot:: http://www.example.com
-    :caption: This is a screenshot for www.example.com
+    :viewport-width: 1280
+    :viewport-height: 960
   ```
 
   You can describe the interaction that you want to have with the webpage
@@ -79,13 +72,6 @@ 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.
 
@@ -96,13 +82,12 @@ class ScreenshotDirective(SphinxDirective):
   """
 
   required_arguments = 1  # URL
-  has_content = True
   option_spec = {
+      **Figure.option_spec,
       'browser': str,
-      'height': directives.positive_int,
-      'width': directives.positive_int,
-      'caption': directives.unchanged,
-      'figclass': directives.unchanged,
+      'viewport-height': directives.positive_int,
+      'viewport-width': directives.positive_int,
+      'interactions': str,
       'pdf': directives.flag,
       'color-scheme': str,
       'full-page': directives.flag,
@@ -112,17 +97,17 @@ class ScreenshotDirective(SphinxDirective):
 
   @staticmethod
   def take_screenshot(
-      url: str, browser_name: str, width: int, height: int, filepath: str,
-      init_script: str, interactions: str, generate_pdf: bool,
+      url: str, browser_name: str, viewport_width: int, viewport_height: int,
+      filepath: str, init_script: str, interactions: str, generate_pdf: bool,
       color_scheme: ColorScheme, full_page: bool,
       context_builder: typing.Optional[typing.Callable[[Browser, str, str],
                                                        BrowserContext]]):
     """Takes a screenshot with Playwright's Chromium browser.
 
     Args:
       url (str): The HTTP/HTTPS URL of the webpage to screenshot.
-      width (int): The width of the screenshot in pixels.
-      height (int): The height of the screenshot in pixels.
+      viewport_width (int): The width of the screenshot in pixels.
+      viewport_height (int): The height of the screenshot in pixels.
       filepath (str): The path to save the screenshot to.
       init_script (str): JavaScript code to be evaluated after the document
         was created but before any of its scripts were run. See more details at
@@ -148,7 +133,10 @@ def take_screenshot(
 
       page = context.new_page()
       page.set_default_timeout(10000)
-      page.set_viewport_size({'width': width, 'height': height})
+      page.set_viewport_size({
+          'width': viewport_width,
+          'height': viewport_height
+      })
 
       try:
         if init_script:
@@ -167,7 +155,10 @@ def take_screenshot(
       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.pdf(
+            width=f'{viewport_width}px',
+            height=f'{viewport_height}px',
+            path=root + '.pdf')
       page.close()
       browser.close()
 
@@ -177,7 +168,7 @@ def evaluate_substitutions(self, text: str) -> str:
       text = text.replace(f"|{key}|", value.astext())
     return text
 
-  def run(self) -> typing.List[nodes.Node]:
+  def run(self) -> typing.Sequence[nodes.Node]:
     screenshot_init_script: str = self.env.config.screenshot_init_script or ''
 
     # Ensure the screenshots directory exists
@@ -187,20 +178,19 @@ def run(self) -> typing.List[nodes.Node]:
     # Parse parameters
     raw_url = self.arguments[0]
     url = self.evaluate_substitutions(raw_url)
+    interactions = self.options.get('interactions', '')
     browser = self.options.get('browser',
                                self.env.config.screenshot_default_browser)
-    height = self.options.get('height',
-                              self.env.config.screenshot_default_height)
-    width = self.options.get('width', self.env.config.screenshot_default_width)
+    viewport_height = self.options.get(
+        'viewport-height', self.env.config.screenshot_default_viewport_height)
+    viewport_width = self.options.get(
+        'viewport-width', self.env.config.screenshot_default_viewport_width)
     color_scheme = self.options.get(
         'color-scheme', self.env.config.screenshot_default_color_scheme)
-    caption_text = self.options.get('caption', '')
-    figclass = self.options.get('figclass', '')
     pdf = 'pdf' in self.options
     full_page = ('full-page' in self.options or
                  self.env.config.screenshot_default_full_page)
     context = self.options.get('context', '')
-    interactions = '\n'.join(self.content)
 
     if urlparse(url).scheme not in {'http', 'https'}:
       raise RuntimeError(
@@ -209,8 +199,8 @@ def run(self) -> typing.List[nodes.Node]:
     # Generate filename based on hash of parameters
     hash_input = "_".join([
         raw_url, browser,
-        str(height),
-        str(width), color_scheme, context, interactions,
+        str(viewport_height),
+        str(viewport_width), color_scheme, context, interactions,
         str(full_page)
     ])
     filename = hashlib.md5(hash_input.encode()).hexdigest() + '.png'
@@ -225,29 +215,18 @@ def run(self) -> typing.List[nodes.Node]:
     # Check if the file already exists. If not, take a screenshot
     if not os.path.exists(filepath):
       fut = self.pool.submit(ScreenshotDirective.take_screenshot, url, browser,
-                             width, height, filepath, screenshot_init_script,
-                             interactions, pdf, color_scheme, full_page,
-                             context_builder)
+                             viewport_width, viewport_height, filepath,
+                             screenshot_init_script, interactions, pdf,
+                             color_scheme, full_page, context_builder)
       fut.result()
 
     # Create image and figure nodes
     docdir = os.path.dirname(self.env.doc2path(self.env.docname))
     rel_ss_dirpath = os.path.relpath(ss_dirpath, start=docdir)
     rel_filepath = os.path.join(rel_ss_dirpath, filename).replace(os.sep, '/')
-    image_node = nodes.image(uri=rel_filepath)
-    figure_node = nodes.figure('', image_node)
-
-    if figclass:
-      figure_node['classes'].append(figclass)
-
-    if caption_text:
-      parsed = nodes.Element()
-      self.state.nested_parse(
-          ViewList([caption_text], source=''), self.content_offset, parsed)
-      figure_node += nodes.caption(parsed[0].source or '', '',
-                                   *parsed[0].children)
 
-    return [figure_node]
+    self.arguments[0] = rel_filepath
+    return super().run()
 
 
 app_threads = {}
@@ -288,12 +267,12 @@ def setup(app: Sphinx) -> Meta:
   app.add_directive('screenshot', ScreenshotDirective)
   app.add_config_value('screenshot_init_script', '', 'env')
   app.add_config_value(
-      'screenshot_default_width',
+      'screenshot_default_viewport_width',
       1280,
       'env',
       description="The default width for screenshots")
   app.add_config_value(
-      'screenshot_default_height',
+      'screenshot_default_viewport_height',
       960,
       'env',
       description="The default height for screenshots")

diff --git a/tests/roots/test-color-schemes/index.rst b/tests/roots/test-color-schemes/index.rst
@@ -1,4 +1,4 @@
 .. screenshot:: |example|
-   :width: 800
-   :height: 600
+   :viewport-width: 800
+   :viewport-height: 600
    :color-scheme: dark
diff --git a/tests/roots/test-default-color-scheme/index.rst b/tests/roots/test-default-color-scheme/index.rst
@@ -1,3 +1,3 @@
 .. screenshot:: |example|
-   :width: 800
-   :height: 600
+   :viewport-width: 800
+   :viewport-height: 600
diff --git a/tests/roots/test-default-size/conf.py b/tests/roots/test-default-size/conf.py
@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 extensions = ['sphinxcontrib.screenshot']
-screenshot_default_width = 1920
-screenshot_default_height = 1200
+screenshot_default_viewport_width = 1920
+screenshot_default_viewport_height = 1200