diff --git a/CHANGELOG.md b/CHANGELOG.md index 81742fb9..c529589b 100644 --- a/CHANGELOG.md +++ b/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! diff --git a/src/py/kaleido/_kaleido_tab/_tab.py b/src/py/kaleido/_kaleido_tab/_tab.py index b3baa24d..8845511b 100644 --- a/src/py/kaleido/_kaleido_tab/_tab.py +++ b/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 = ""): @@ -100,6 +105,8 @@ async def navigate(self, url: str | Path = ""): # requires a couple extra lines self.js_logger.reset() + await self._apply_headers() + # reload is truly so close to navigate async def reload(self): """Reload the tab, and set the javascript runtime id.""" @@ -118,6 +125,20 @@ async def reload(self): self.js_logger.reset() + await self._apply_headers() + + async def _apply_headers(self): + """Apply extra HTTP headers to the tab if configured.""" + if self._headers: + _logger.debug2(f"Setting extra HTTP headers on {self.tab}") + _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, diff --git a/src/py/kaleido/kaleido.py b/src/py/kaleido/kaleido.py index 8635c202..f0771f3c 100644 --- a/src/py/kaleido/kaleido.py +++ b/src/py/kaleido/kaleido.py @@ -105,7 +105,7 @@ 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, @@ -113,6 +113,7 @@ def __init__( 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)) diff --git a/src/py/pyproject.toml b/src/py/pyproject.toml index 47dd838c..fe0903d9 100644 --- a/src/py/pyproject.toml +++ b/src/py/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "typing-extensions>=4.12.2", "hypothesis>=6.113.0", "pyright>=1.1.406", + "ruff", ] pickles = [ diff --git a/src/py/tests/test_kaleido.py b/src/py/tests/test_kaleido.py index a4faaafa..a2f213f6 100644 --- a/src/py/tests/test_kaleido.py +++ b/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" diff --git a/src/py/uv.lock b/src/py/uv.lock index c1c07cfe..0e2abb1d 100644 --- a/src/py/uv.lock +++ b/src/py/uv.lock @@ -507,6 +507,7 @@ dev = [ { name = "pytest-timeout" }, { name = "pytest-xdist", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, { name = "pytest-xdist", version = "3.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "ruff" }, { name = "typing-extensions", version = "4.13.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.9'" }, ] @@ -544,6 +545,7 @@ dev = [ { name = "pytest-order", specifier = ">=1.3.0" }, { name = "pytest-timeout", specifier = ">=2.4.0" }, { name = "pytest-xdist" }, + { name = "ruff" }, { name = "typing-extensions", specifier = ">=4.12.2" }, ] pickles = [ @@ -2180,6 +2182,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, ] +[[package]] +name = "ruff" +version = "0.15.11" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e4/8d/192f3d7103816158dfd5ea50d098ef2aec19194e6cbccd4b3485bdb2eb2d/ruff-0.15.11.tar.gz", hash = "sha256:f092b21708bf0e7437ce9ada249dfe688ff9a0954fc94abab05dcea7dcd29c33", size = 4637264, upload-time = "2026-04-16T18:46:26.58Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/1e/6aca3427f751295ab011828e15e9bf452200ac74484f1db4be0197b8170b/ruff-0.15.11-py3-none-linux_armv6l.whl", hash = "sha256:e927cfff503135c558eb581a0c9792264aae9507904eb27809cdcff2f2c847b7", size = 10607943, upload-time = "2026-04-16T18:46:05.967Z" }, + { url = "https://files.pythonhosted.org/packages/e7/26/1341c262e74f36d4e84f3d6f4df0ac68cd53331a66bfc5080daa17c84c0b/ruff-0.15.11-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:7a1b5b2938d8f890b76084d4fa843604d787a912541eae85fd7e233398bbb73e", size = 10988592, upload-time = "2026-04-16T18:46:00.742Z" }, + { url = "https://files.pythonhosted.org/packages/03/71/850b1d6ffa9564fbb6740429bad53df1094082fe515c8c1e74b6d8d05f18/ruff-0.15.11-py3-none-macosx_11_0_arm64.whl", hash = "sha256:d4176f3d194afbdaee6e41b9ccb1a2c287dba8700047df474abfbe773825d1cb", size = 10338501, upload-time = "2026-04-16T18:46:03.723Z" }, + { url = "https://files.pythonhosted.org/packages/f2/11/cc1284d3e298c45a817a6aadb6c3e1d70b45c9b36d8d9cce3387b495a03a/ruff-0.15.11-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3b17c886fb88203ced3afe7f14e8d5ae96e9d2f4ccc0ee66aa19f2c2675a27e4", size = 10670693, upload-time = "2026-04-16T18:46:41.941Z" }, + { url = "https://files.pythonhosted.org/packages/ce/9e/f8288b034ab72b371513c13f9a41d9ba3effac54e24bfb467b007daee2ca/ruff-0.15.11-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:49fafa220220afe7758a487b048de4c8f9f767f37dfefad46b9dd06759d003eb", size = 10416177, upload-time = "2026-04-16T18:46:21.717Z" }, + { url = "https://files.pythonhosted.org/packages/85/71/504d79abfd3d92532ba6bbe3d1c19fada03e494332a59e37c7c2dabae427/ruff-0.15.11-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f2ab8427e74a00d93b8bda1307b1e60970d40f304af38bccb218e056c220120d", size = 11221886, upload-time = "2026-04-16T18:46:15.086Z" }, + { url = "https://files.pythonhosted.org/packages/43/5a/947e6ab7a5ad603d65b474be15a4cbc6d29832db5d762cd142e4e3a74164/ruff-0.15.11-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:195072c0c8e1fc8f940652073df082e37a5d9cb43b4ab1e4d0566ab8977a13b7", size = 12075183, upload-time = "2026-04-16T18:46:07.944Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a1/0b7bb6268775fdd3a0818aee8efd8f5b4e231d24dd4d528ced2534023182/ruff-0.15.11-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a3a0996d486af3920dec930a2e7daed4847dfc12649b537a9335585ada163e9e", size = 11516575, upload-time = "2026-04-16T18:46:31.687Z" }, + { url = "https://files.pythonhosted.org/packages/30/c3/bb5168fc4d233cc06e95f482770d0f3c87945a0cd9f614b90ea8dc2f2833/ruff-0.15.11-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bef2cb556d509259f1fe440bb9cd33c756222cf0a7afe90d15edf0866702431", size = 11306537, upload-time = "2026-04-16T18:46:36.988Z" }, + { url = "https://files.pythonhosted.org/packages/e4/92/4cfae6441f3967317946f3b788136eecf093729b94d6561f963ed810c82e/ruff-0.15.11-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:030d921a836d7d4a12cf6e8d984a88b66094ccb0e0f17ddd55067c331191bf19", size = 11296813, upload-time = "2026-04-16T18:46:24.182Z" }, + { url = "https://files.pythonhosted.org/packages/43/26/972784c5dde8313acde8ac71ba8ac65475b85db4a2352a76c9934361f9bc/ruff-0.15.11-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:0e783b599b4577788dbbb66b9addcef87e9a8832f4ce0c19e34bf55543a2f890", size = 10633136, upload-time = "2026-04-16T18:46:39.802Z" }, + { url = "https://files.pythonhosted.org/packages/5b/53/3985a4f185020c2f367f2e08a103032e12564829742a1b417980ce1514a0/ruff-0.15.11-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:ae90592246625ba4a34349d68ec28d4400d75182b71baa196ddb9f82db025ef5", size = 10424701, upload-time = "2026-04-16T18:46:10.381Z" }, + { url = "https://files.pythonhosted.org/packages/d3/57/bf0dfb32241b56c83bb663a826133da4bf17f682ba8c096973065f6e6a68/ruff-0.15.11-py3-none-musllinux_1_2_i686.whl", hash = "sha256:1f111d62e3c983ed20e0ca2e800f8d77433a5b1161947df99a5c2a3fb60514f0", size = 10873887, upload-time = "2026-04-16T18:46:29.157Z" }, + { url = "https://files.pythonhosted.org/packages/02/05/e48076b2a57dc33ee8c7a957296f97c744ca891a8ffb4ffb1aaa3b3f517d/ruff-0.15.11-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:06f483d6646f59eaffba9ae30956370d3a886625f511a3108994000480621d1c", size = 11404316, upload-time = "2026-04-16T18:46:19.462Z" }, + { url = "https://files.pythonhosted.org/packages/88/27/0195d15fe7a897cbcba0904792c4b7c9fdd958456c3a17d2ea6093716a9a/ruff-0.15.11-py3-none-win32.whl", hash = "sha256:476a2aa56b7da0b73a3ee80b6b2f0e19cce544245479adde7baa65466664d5f3", size = 10655535, upload-time = "2026-04-16T18:46:12.47Z" }, + { url = "https://files.pythonhosted.org/packages/3a/5e/c927b325bd4c1d3620211a4b96f47864633199feed60fa936025ab27e090/ruff-0.15.11-py3-none-win_amd64.whl", hash = "sha256:8b6756d88d7e234fb0c98c91511aae3cd519d5e3ed271cae31b20f39cb2a12a3", size = 11779692, upload-time = "2026-04-16T18:46:17.268Z" }, + { url = "https://files.pythonhosted.org/packages/63/b6/aeadee5443e49baa2facd51131159fd6301cc4ccfc1541e4df7b021c37dd/ruff-0.15.11-py3-none-win_arm64.whl", hash = "sha256:063fed18cc1bbe0ee7393957284a6fe8b588c6a406a285af3ee3f46da2391ee4", size = 11032614, upload-time = "2026-04-16T18:46:34.487Z" }, +] + [[package]] name = "scipy" version = "1.10.1"