Handle pathlib.Path in write_image

Author: maresb

Diff for: packages/python/plotly/plotly/io/_kaleido.py

@@ -2,6 +2,7 @@
 from six import string_types
 import os
 import json
+from pathlib import Path, PurePath
 import plotly
 from plotly.io._utils import validate_coerce_fig_to_dict
 
@@ -169,7 +170,7 @@ def write_image(
 
     file: str or writeable
         A string representing a local file path or a writeable object
-        (e.g. an open file descriptor)
+        (e.g. a pathlib.Path object or an open file descriptor)
 
     format: str or None
         The desired image format. One of
@@ -228,14 +229,23 @@ def write_image(
     -------
     None
     """
-    # Check if file is a string
-    # -------------------------
-    file_is_str = isinstance(file, string_types)
+    # Try to cast `file` as a pathlib object `path`.
+    # ----------------------------------------------
+    if isinstance(file, string_types):
+        # Use the standard pathlib constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, PurePath):  # PurePath is the most general pathlib object.
+        # `file` is already a pathlib object.
+        path = file
+    else:
+        # We could not make a pathlib object out of file. Either `file` is an open file
+        # descriptor with a `write()` method or it's an invalid object.
+        path = None
 
     # Infer format if not specified
     # -----------------------------
-    if file_is_str and format is None:
-        _, ext = os.path.splitext(file)
+    if path is not None and format is None:
+        ext = path.suffix
         if ext:
             format = ext.lstrip(".")
         else:
@@ -267,11 +277,25 @@ def write_image(
 
     # Open file
     # ---------
-    if file_is_str:
-        with open(file, "wb") as f:
-            f.write(img_data)
+    if path is None:
+        # We previously failed to make sense of `file` as a pathlib object.
+        # Attempt to write to `file` as an open file descriptor.
+        try:
+            file.write(img_data)
+            return
+        except AttributeError:
+            pass
+        raise ValueError(
+            """
+The 'file' argument '{file}' is not a string, pathlib.Path object, or file descriptor.
+""".format(
+                file=file
+            )
+        )
     else:
-        file.write(img_data)
+        # We previously succeeded in interpreting `file` as a pathlib object.
+        # Now we can use `write_bytes()`.
+        path.write_bytes(img_data)
 
 
 def full_figure_for_development(fig, warn=True, as_dict=False):

Diff for: packages/python/plotly/plotly/tests/test_optional/test_kaleido/test_kaleido.py

@@ -1,16 +1,56 @@
 import plotly.io as pio
 import plotly.io.kaleido
-import sys
 from contextlib import contextmanager
-
-if sys.version_info >= (3, 3):
-    from unittest.mock import Mock
-else:
-    from mock import Mock
+from io import BytesIO
+from pathlib import Path
+from unittest.mock import Mock
 
 fig = {"layout": {"title": {"text": "figure title"}}}
 
 
+def make_writeable_mocks():
+    """Produce some mocks which we will use for testing the `write_image()` function.
+
+    These mocks should be passed as the `file=` argument to `write_image()`.
+
+    The tests should verify that the method specified in the `active_write_function`
+    attribute is called once, and that scope.transform is called with the `format=`
+    argument specified by the `.expected_format` attribute.
+
+    In total we provide two mocks: one for a writable file descriptor, and other for a
+    pathlib.Path object.
+    """
+
+    # Part 1: A mock for a file descriptor
+    # ------------------------------------
+    mock_file_descriptor = Mock()
+
+    # A file descriptor has no write_bytes method, unlike a pathlib Path.
+    del mock_file_descriptor.write_bytes
+
+    # The expected write method for a file descriptor is .write
+    mock_file_descriptor.active_write_function = mock_file_descriptor.write
+
+    # Since there is no filename, there should be no format detected.
+    mock_file_descriptor.expected_format = None
+
+    # Part 2: A mock for a pathlib path
+    # ---------------------------------
+    mock_pathlib_path = Mock(spec=Path)
+
+    # A pathlib Path object has no write method, unlike a file descriptor.
+    del mock_pathlib_path.write
+
+    # The expected write method for a pathlib Path is .write_bytes
+    mock_pathlib_path.active_write_function = mock_pathlib_path.write_bytes
+
+    # Mock a path with PNG suffix
+    mock_pathlib_path.suffix = ".png"
+    mock_pathlib_path.expected_format = "png"
+
+    return mock_file_descriptor, mock_pathlib_path
+
+
 @contextmanager
 def mocked_scope():
     # Code to acquire resource, e.g.:
@@ -44,15 +84,19 @@ def test_kaleido_engine_to_image():
 
 
 def test_kaleido_engine_write_image():
-    writeable_mock = Mock()
-    with mocked_scope() as scope:
-        pio.write_image(fig, writeable_mock, engine="kaleido", validate=False)
+    for writeable_mock in make_writeable_mocks():
+        with mocked_scope() as scope:
+            pio.write_image(fig, writeable_mock, engine="kaleido", validate=False)
 
-    scope.transform.assert_called_with(
-        fig, format=None, width=None, height=None, scale=None
-    )
+        scope.transform.assert_called_with(
+            fig,
+            format=writeable_mock.expected_format,
+            width=None,
+            height=None,
+            scale=None,
+        )
 
-    assert writeable_mock.write.call_count == 1
+        assert writeable_mock.active_write_function.call_count == 1
 
 
 def test_kaleido_engine_to_image_kwargs():
@@ -73,24 +117,24 @@ def test_kaleido_engine_to_image_kwargs():
 
 
 def test_kaleido_engine_write_image_kwargs():
-    writeable_mock = Mock()
-    with mocked_scope() as scope:
-        pio.write_image(
-            fig,
-            writeable_mock,
-            format="jpg",
-            width=700,
-            height=600,
-            scale=2,
-            engine="kaleido",
-            validate=False,
+    for writeable_mock in make_writeable_mocks():
+        with mocked_scope() as scope:
+            pio.write_image(
+                fig,
+                writeable_mock,
+                format="jpg",
+                width=700,
+                height=600,
+                scale=2,
+                engine="kaleido",
+                validate=False,
+            )
+
+        scope.transform.assert_called_with(
+            fig, format="jpg", width=700, height=600, scale=2
         )
 
-    scope.transform.assert_called_with(
-        fig, format="jpg", width=700, height=600, scale=2
-    )
-
-    assert writeable_mock.write.call_count == 1
+        assert writeable_mock.active_write_function.call_count == 1
 
 
 def test_image_renderer():
@@ -105,3 +149,17 @@ def test_image_renderer():
         height=renderer.height,
         scale=renderer.scale,
     )
+
+
+def test_bytesio():
+    """Verify that writing to a BytesIO object contains the same data as to_image().
+
+    The goal of this test is to ensure that Plotly correctly handles a writable buffer
+    which doesn't correspond to a filesystem path.
+    """
+    bio = BytesIO()
+    pio.write_image(fig, bio, format="jpg", engine="kaleido", validate=False)
+    bio.seek(0)  # Rewind to the beginning of the buffer, otherwise read() returns b''.
+    bio_bytes = bio.read()
+    to_image_bytes = pio.to_image(fig, format="jpg", engine="kaleido", validate=False)
+    assert bio_bytes == to_image_bytes

Diff for: packages/python/plotly/plotly/tests/test_orca/test_to_image.py

@@ -6,6 +6,7 @@
 import pytest
 import sys
 import pandas as pd
+from io import BytesIO
 
 if sys.version_info >= (3, 3):
     from unittest.mock import MagicMock
@@ -234,7 +235,12 @@ def test_write_image_writeable(fig1, format):
         fig1, mock_file, format=format, width=700, height=500, engine="orca"
     )
 
-    mock_file.write.assert_called_once_with(expected_bytes)
+    if mock_file.write_bytes.called:
+        mock_file.write_bytes.assert_called_once_with(expected_bytes)
+    elif mock_file.write.called:
+        mock_file.write.assert_called_once_with(expected_bytes)
+    else:
+        assert "Neither write nor write_bytes was called."
 
 
 def test_write_image_string_format_inference(fig1, format):
@@ -347,3 +353,17 @@ def test_invalid_figure_json():
     )
 
     assert "400: invalid or malformed request syntax" in str(err.value)
+
+
+def test_bytesio(fig1):
+    """Verify that writing to a BytesIO object contains the same data as to_image().
+
+    The goal of this test is to ensure that Plotly correctly handles a writable buffer
+    which doesn't correspond to a filesystem path.
+    """
+    bio = BytesIO()
+    pio.write_image(fig1, bio, format="jpg", validate=False)
+    bio.seek(0)
+    bio_bytes = bio.read()
+    to_image_bytes = pio.to_image(fig1, format="jpg", validate=False)
+    assert bio_bytes == to_image_bytes