refactor write_images() to use kaleido.write_fig_from_object_sync()

Author: emilykl

plotly/io/_kaleido.py

@@ -6,11 +6,17 @@
 import warnings
 
 import plotly
-from plotly.io._utils import validate_coerce_fig_to_dict, as_individual_args
+from plotly.io._utils import validate_coerce_fig_to_dict, broadcast_args_to_dicts
 from plotly.io import defaults
 
 ENGINE_SUPPORT_TIMELINE = "September 2025"
 
+PLOTLY_GET_CHROME_ERROR_MSG = """
+
+Kaleido requires Google Chrome to be installed. Install it by running:
+    $ plotly_get_chrome
+"""
+
 
 # TODO: Remove --pre flag once Kaleido v1 full release is available
 KALEIDO_DEPRECATION_MSG = f"""
@@ -158,6 +164,7 @@ def as_path_object(file: str | Path) -> Path | None:
         path = None
     return path
 
+
 def infer_format(path: Path | None, format: str | None) -> str | None:
     if path is not None and format is None:
         ext = path.suffix
@@ -177,7 +184,6 @@ def infer_format(path: Path | None, format: str | None) -> str | None:
     return format
 
 
-
 def to_image(
     fig,
     format=None,
@@ -331,13 +337,7 @@ def to_image(
                 ),
             )
         except choreographer.errors.ChromeNotFoundError:
-            raise RuntimeError(
-                """
-
-Kaleido requires Google Chrome to be installed. Install it by running:
-    $ plotly_get_chrome
-"""
-            )
+            raise RuntimeError(PLOTLY_GET_CHROME_ERROR_MSG)
 
     else:
         # Kaleido v0
@@ -481,7 +481,7 @@ def write_image(
 
 
 def write_images(
-    figs,
+    fig,
     file,
     format=None,
     scale=None,
@@ -494,13 +494,17 @@ def write_images(
     calling write_image() multiple times. This function can only be used with the Kaleido
     engine, v1.0.0 or greater.
 
+    This function accepts the same arguments as write_image() (minus the `engine` argument),
+    except that any of the arguments may be either a single value or an iterable of values.
+    If multiple arguments are iterable, they must all have the same length.
+
     Parameters
     ----------
-    figs:
+    fig:
         Iterable of figure objects or dicts representing a figure
 
-    directory: str or writeable
-        A string or pathlib.Path object representing a local directory path.
+    file: str or writeable
+        Iterables of strings or pathlib.Path objects representing local file paths to write to.
 
     format: str or None
         The desired image format. One of
@@ -562,39 +566,48 @@ def write_images(
 """
         )
 
-    # Try to cast `file` as a pathlib object `path`.
-    path = as_path_object(file)
-
-    # Infer image format if not specified
-    format = infer_format(path, format)
-
-    # Convert figures to dicts (and validate if requested)
-    # TODO: Keep same iterable type
-    fig_dicts = [validate_coerce_fig_to_dict(fig, validate) for fig in figs]    
-
-    kaleido.write_fig_sync(
-        fig_dicts,
-        directory=path,
-        opts=dict(
-            format=format or defaults.default_format,
-            width=width or defaults.default_width,
-            height=height or defaults.default_height,
-            scale=scale or defaults.default_scale,
-        ),
+    # Broadcast arguments into correct format for passing to Kaleido
+    arg_dicts = broadcast_args_to_dicts(
+        fig=fig,
+        file=file,
+        format=format,
+        scale=scale,
+        width=width,
+        height=height,
+        validate=validate,
     )
 
-    # # Get individual arguments
-    # individual_args, individual_kwargs = as_individual_args(*args, **kwargs)
-
-    # if kaleido_available() and kaleido_major() > 0:
-    #     # Kaleido v1
-    #     # TODO: Use a single shared kaleido instance for all images
-    #     for a, kw in zip(individual_args, individual_kwargs):
-    #         write_image(*a, **kw)
-    # else:
-    #     # Kaleido v0, or orca
-    #     for a, kw in zip(individual_args, individual_kwargs):
-    #         write_image(*a, **kw)
+    # For each dict:
+    #   - convert figures to dicts (and validate if requested)
+    #   - try to cast `file` as a Path object
+    for d in arg_dicts:
+        d["fig"] = validate_coerce_fig_to_dict(d["fig"], d["validate"])
+        d["file"] = as_path_object(d["file"])
+
+    # Reshape arg_dicts into correct format for passing to Kaleido
+    # We call infer_format() here rather than above so that the `file` argument
+    # has already been cast to a Path object.
+    # Also insert defaults for any missing arguments as needed
+    kaleido_specs = [
+        {
+            "fig": d["fig"],
+            "path": d["file"],
+            "opts": dict(
+                format=infer_format(d["file"], d["format"]) or defaults.default_format,
+                width=d["width"] or defaults.default_width,
+                height=d["height"] or defaults.default_height,
+                scale=d["scale"] or defaults.default_scale,
+            ),
+        }
+        for d in arg_dicts
+    ]
+
+    import choreographer
+
+    try:
+        kaleido.write_fig_from_object_sync(kaleido_specs)
+    except choreographer.errors.ChromeNotFoundError:
+        raise RuntimeError(PLOTLY_GET_CHROME_ERROR_MSG)
 
 
 def full_figure_for_development(fig, warn=True, as_dict=False):

plotly/io/_utils.py

@@ -43,55 +43,47 @@ def validate_coerce_output_type(output_type):
     return cls
 
 
-def as_individual_args(*args, **kwargs):
+def broadcast_args_to_dicts(**kwargs):
     """
-    Given one or more positional or keyword arguments which may be either a single value
-    or a list of values, return a list of lists and a list of dictionaries
-    by expanding the single values into lists.
+    Given one or more keyword arguments which may be either a single value or a list of values,
+    return a list of keyword dictionaries by broadcasting the single valuesacross all the dicts.
     If more than one item in the input is a list, all lists must be the same length.
 
     Parameters
     ----------
-    *args: list
-        The positional arguments
     **kwargs: dict
         The keyword arguments
 
     Returns
     -------
-    list of lists
-        A list of lists
     list of dicts
         A list of dictionaries
+
+    Raises
+    ------
+    ValueError
+        If any of the input lists are not the same length
     """
     # Check that all list arguments have the same length,
     # and find out what that length is
     # If there are no list arguments, length is 1
-    list_lengths = [
-        len(v) for v in args + tuple(kwargs.values()) if isinstance(v, list)
-    ]
+    list_lengths = [len(v) for v in tuple(kwargs.values()) if isinstance(v, list)]
     if list_lengths and len(set(list_lengths)) > 1:
         raise ValueError("All list arguments must have the same length.")
     list_length = list_lengths[0] if list_lengths else 1
 
     # Expand all arguments to lists of the same length
-    expanded_args = [[v] * list_length if not isinstance(v, list) else v for v in args]
     expanded_kwargs = {
         k: [v] * list_length if not isinstance(v, list) else v
         for k, v in kwargs.items()
     }
-
-    # Reshape into a list of lists
-    # Each list represents the positional arguments for a single function call
-    list_of_args = [[v[i] for v in expanded_args] for i in range(list_length)]
-
     # Reshape into a list of dictionaries
     # Each dictionary represents the keyword arguments for a single function call
     list_of_kwargs = [
         {k: v[i] for k, v in expanded_kwargs.items()} for i in range(list_length)
     ]
 
-    return list_of_args, list_of_kwargs
+    return list_of_kwargs
 
 
 def plotly_cdn_url(cdn_ver=get_plotlyjs_version()):