diff --git a/plotly/io/__init__.py b/plotly/io/__init__.py index 87f9c3a49a..e2d68242da 100644 --- a/plotly/io/__init__.py +++ b/plotly/io/__init__.py @@ -17,7 +17,7 @@ from ._html import to_html, write_html from ._renderers import renderers, show from . import base_renderers - from ._kaleido import defaults + from ._kaleido import defaults, get_chrome __all__ = [ "to_image", @@ -38,6 +38,7 @@ "base_renderers", "full_figure_for_development", "defaults", + "get_chrome", ] else: __all__, __getattr__, __dir__ = relative_import( @@ -59,6 +60,7 @@ "._renderers.renderers", "._renderers.show", "._kaleido.defaults", + "._kaleido.get_chrome", ], ) diff --git a/plotly/io/_kaleido.py b/plotly/io/_kaleido.py index 6775c333f1..334227076f 100644 --- a/plotly/io/_kaleido.py +++ b/plotly/io/_kaleido.py @@ -775,11 +775,13 @@ def full_figure_for_development( return go.Figure(fig, skip_invalid=True) -def get_chrome() -> None: +def plotly_get_chrome() -> None: """ Install Google Chrome for Kaleido (Required for Plotly image export). - This function can be run from the command line using the command `plotly_get_chrome` - defined in pyproject.toml + This function is a command-line wrapper for `plotly.io.get_chrome()`. + + When running from the command line, use the command `plotly_get_chrome`; + when calling from Python code, use `plotly.io.get_chrome()`. """ usage = """ @@ -813,7 +815,6 @@ def get_chrome() -> None: # Handle "--path" flag chrome_install_path = None - user_specified_path = False if "--path" in cli_args: path_index = cli_args.index("--path") + 1 if path_index < len(cli_args): @@ -821,8 +822,53 @@ def get_chrome() -> None: cli_args.remove("--path") cli_args.remove(chrome_install_path) chrome_install_path = Path(chrome_install_path) - user_specified_path = True + + # If any arguments remain, command syntax was incorrect -- print usage and exit + if len(cli_args) > 1: + print(usage) + sys.exit(1) + + if not cli_yes: + print( + f""" +Plotly will install a copy of Google Chrome to be used for generating static images of plots. +Chrome will be installed at: {chrome_install_path}""" + ) + response = input("Do you want to proceed? [y/n] ") + if not response or response[0].lower() != "y": + print("Cancelled") + return + print("Installing Chrome for Plotly...") + exe_path = get_chrome(chrome_install_path) + print("Chrome installed successfully.") + print(f"The Chrome executable is now located at: {exe_path}") + + +def get_chrome(path: Union[str, Path, None] = None) -> Path: + """ + Get the path to the Chrome executable for Kaleido. + This function is used by the `plotly_get_chrome` command line utility. + + Parameters + ---------- + path: str or Path or None + The path to the directory where Chrome should be installed. + If None, the default download path will be used. + """ + if not kaleido_available() or kaleido_major() < 1: + raise ValueError( + """ +This command requires Kaleido v1.0.0 or greater. +Install it using `pip install 'kaleido>=1.0.0'` or `pip install 'plotly[kaleido]'`." +""" + ) + + # Use default download path if no path was specified + if path: + user_specified_path = True + chrome_install_path = Path(path) # Ensure it's a Path object else: + user_specified_path = False from choreographer.cli.defaults import default_download_path chrome_install_path = default_download_path @@ -848,25 +894,7 @@ def get_chrome() -> None: """ ) - # If any arguments remain, command syntax was incorrect -- print usage and exit - if len(cli_args) > 1: - print(usage) - sys.exit(1) - - if not cli_yes: - print( - f""" -Plotly will install a copy of Google Chrome to be used for generating static images of plots. -Chrome will be installed at: {chrome_install_path}""" - ) - response = input("Do you want to proceed? [y/n] ") - if not response or response[0].lower() != "y": - print("Cancelled") - return - print("Installing Chrome for Plotly...") - exe_path = kaleido.get_chrome_sync(path=chrome_install_path) - print("Chrome installed successfully.") - print(f"The Chrome executable is now located at: {exe_path}") + return kaleido.get_chrome_sync(path=chrome_install_path) __all__ = ["to_image", "write_image", "scope", "full_figure_for_development"] diff --git a/pyproject.toml b/pyproject.toml index 49a821cfb2..9d060de62b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,7 @@ dev = [ ] [project.scripts] -plotly_get_chrome = "plotly.io._kaleido:get_chrome" +plotly_get_chrome = "plotly.io._kaleido:plotly_get_chrome" [tool.pytest.ini_options] markers = [ diff --git a/tests/test_optional/test_kaleido/test_kaleido.py b/tests/test_optional/test_kaleido/test_kaleido.py index 47ac263ccf..7e9c09a369 100644 --- a/tests/test_optional/test_kaleido/test_kaleido.py +++ b/tests/test_optional/test_kaleido/test_kaleido.py @@ -293,3 +293,24 @@ def test_fig_to_image(): mock_calc_fig.assert_called_once() args, _ = mock_calc_fig.call_args assert args[0] == test_fig.to_dict() + + +def test_get_chrome(): + """Test that plotly.io.get_chrome() can be called.""" + + if not kaleido_available() or kaleido_major() < 1: + # Test that ValueError is raised when Kaleido requirements aren't met + with pytest.raises( + ValueError, match="This command requires Kaleido v1.0.0 or greater" + ): + pio.get_chrome() + else: + # Test normal operation when Kaleido v1+ is available + with patch( + "plotly.io._kaleido.kaleido.get_chrome_sync", + return_value="/mock/path/to/chrome", + ) as mock_get_chrome: + pio.get_chrome() + + # Verify that kaleido.get_chrome_sync was called + mock_get_chrome.assert_called_once()