Merge pull request #5282 from plotly/5272-add-plotly-io-get-chrome

Expose `plotly.io.get_chrome()`

Author: emilykl

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",
         ],
     )
 

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,16 +815,60 @@ 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):
             chrome_install_path = cli_args[path_index]
             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"]

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 = [

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()