Merge pull request #446 from plotly/cam/445/add-headers

feat: Allow for request headers to be passed in

Author: camdecoster

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## Unreleased
 
+### Added
+- Allow for request headers to be added to Choreographer calls [[#443](https://github.com/plotly/Kaleido/issues/443)]
+
 ### Fixed
 - Fix issue where exporting large figures could cause hang [[#442](https://github.com/plotly/Kaleido/pull/442)], with thanks to @EliasTalcott for the contribution!
 

src/py/kaleido/_kaleido_tab/_tab.py

@@ -53,15 +53,20 @@ class _KaleidoTab:
     js_logger: _js_logger.JavascriptLogger
     """A log for recording javascript."""
 
-    def __init__(self, tab):
+    def __init__(self, tab, *, headers: dict[str, str] | None = None):
         """
         Create a new _KaleidoTab.
 
         Args:
             tab: the choreographer tab to wrap.
 
+            headers (dict[str, str] | None, optional):
+                Extra HTTP headers to send with every request made by the
+                browser tab. Defaults to None.
+
         """
         self.tab = tab
+        self._headers = headers
         self.js_logger = _js_logger.JavascriptLogger(self.tab)
 
     async def navigate(self, url: str | Path = ""):
@@ -79,6 +84,9 @@ async def navigate(self, url: str | Path = ""):
         # Subscribe to event indicating page ready.
         page_ready = _subscribe_new(self.tab, "Page.loadEventFired")
 
+        # Apply headers if they exist
+        await self._apply_headers()
+
         # Navigating page. This will trigger the above events.
         _logger.debug2(f"Calling Page.navigate on {self.tab}")
         _raise_error(await self.tab.send_command("Page.navigate", params={"url": url}))
@@ -118,6 +126,19 @@ async def reload(self):
 
         self.js_logger.reset()
 
+    async def _apply_headers(self):
+        """Apply extra HTTP headers to the tab if configured."""
+        if self._headers:
+            _logger.debug(f"Setting extra HTTP headers on {self.tab}")
+            _logger.debug2(f"Extra headers are: {self._headers}")
+            _raise_error(await self.tab.send_command("Network.enable"))
+            _raise_error(
+                await self.tab.send_command(
+                    "Network.setExtraHTTPHeaders",
+                    params={"headers": self._headers},
+                )
+            )
+
     async def _calc_fig(
         self,
         spec: fig_tools.Spec,

src/py/kaleido/kaleido.py

@@ -105,14 +105,15 @@ class Kaleido(choreo.Browser):
 
     ### KALEIDO LIFECYCLE FUNCTIONS ###
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         # *args: Any, force named vars for all choreographer passthrough
         n: int = 1,
         timeout: float | None = 90,
         page_generator: None | PageGenerator | str | Path = None,
         plotlyjs: str | Path | None = None,
         mathjax: str | Path | Literal[False] | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -145,6 +146,12 @@ def __init__(
                 disabled. Defaults to None- which means to use version 2.35 via
                 CDN.
 
+            headers (dict[str, str] | None, optional):
+                A dictionary of extra HTTP headers to send with every request
+                made by the browser (e.g. {"Referer": "https://example.com/"}).
+                Uses the Chrome DevTools Protocol Network.setExtraHTTPHeaders.
+                Defaults to None.
+
             **kwargs (Any):
                 Additional keyword arguments passed through to the underlying
                 Choreographer.browser constructor. Notable options include
@@ -172,6 +179,7 @@ def __init__(
         self._n = n
         self._plotlyjs = plotlyjs
         self._mathjax = mathjax
+        self._headers = headers
 
         # Diagnostic
         _logger.debug(f"Timeout: {self._timeout}")
@@ -229,7 +237,7 @@ async def _conform_tabs(self, tabs: Listish[choreo.Tab] | None = None) -> None:
             _logger.debug2(f"Subscribing * to tab: {tab}.")
             tab.subscribe("*", _utils.event_printer(f"tab-{i!s}: Event Dump:"))
 
-        kaleido_tabs = [_KaleidoTab(tab) for tab in tabs]
+        kaleido_tabs = [_KaleidoTab(tab, headers=self._headers) for tab in tabs]
 
         await asyncio.gather(*(tab.navigate(self._index) for tab in kaleido_tabs))
 

src/py/pyproject.toml

@@ -55,6 +55,7 @@ dev = [
   "typing-extensions>=4.12.2",
   "hypothesis>=6.113.0",
   "pyright>=1.1.406",
+  "ruff",
 ]
 
 pickles = [

src/py/tests/test_kaleido.py

@@ -509,3 +509,38 @@ async def test_plotlyjs_mathjax_injection(plotlyjs, mathjax):
         finally:
             # Put the tab back in the queue
             await k.tabs_ready.put(tab)
+
+
+async def test_headers_stored_on_tabs():
+    """Test that custom headers are passed through to _KaleidoTab instances."""
+    test_headers = {"Referer": "https://example.com/", "X-Custom": "value"}
+
+    async with Kaleido(headers=test_headers, n=1) as k:
+        tab = await k.tabs_ready.get()
+        try:
+            assert tab._headers == test_headers  # noqa: SLF001
+        finally:
+            await k.tabs_ready.put(tab)
+
+
+async def test_headers_none_by_default():
+    """Test that headers default to None when not specified."""
+    async with Kaleido(n=1) as k:
+        tab = await k.tabs_ready.get()
+        try:
+            assert tab._headers is None  # noqa: SLF001
+        finally:
+            await k.tabs_ready.put(tab)
+
+
+async def test_headers_rendering_works(simple_figure_with_bytes):
+    """Test that rendering still works correctly when headers are set."""
+    test_headers = {"Referer": "https://example.com/"}
+
+    async with Kaleido(headers=test_headers) as k:
+        result = await k.calc_fig(
+            simple_figure_with_bytes["fig"],
+            opts=simple_figure_with_bytes["opts"],
+        )
+
+    assert result[:8] == b"\x89PNG\r\n\x1a\n", "Generated data is not a valid PNG"