Rework logic to actually use pathlib

Author: maresb

packages/python/plotly/plotly/io/_kaleido.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 from six import string_types
 import os
+from pathlib import Path, PurePath
 import json
 import plotly
 from plotly.io._utils import validate_coerce_fig_to_dict
@@ -215,14 +216,19 @@ def write_image(
     -------
     None
     """
-    # Check if file is a string
-    # -------------------------
-    file_is_str = isinstance(file, string_types)
+    # Check if file can be interpreted as a pathlib object
+    # ----------------------------------------------------
+    if isinstance(file, string_types):
+        path = Path(file)
+    elif isinstance(file, PurePath):
+        path = file
+    else:
+        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:
@@ -254,26 +260,16 @@ def write_image(
 
     # Open file
     # ---------
-    if file_is_str:
-        with open(file, "wb") as f:
-            f.write(img_data)
+    if path is not None:
+        file.write_bytes(img_data)
     else:
-        # Handle a pathlib.Path object
-        # ----------------------------
-        try:
-            file.write_bytes(img_data)
-            return
-        except AttributeError:
-            pass
         # Handle an open file descriptor
         # ------------------------------
         try:
             file.write(img_data)
             return
         except AttributeError:
             pass
-        # Handlers failed
-        # ---------------
         raise ValueError(
             """
 The 'file' argument '{file}' is not a string, pathlib.Path object, or file descriptor.

packages/python/plotly/plotly/tests/test_optional/test_kaleido/test_kaleido.py

@@ -1,6 +1,7 @@
 import plotly.io as pio
 import plotly.io.kaleido
 import sys
+from pathlib import Path
 from contextlib import contextmanager
 
 if sys.version_info >= (3, 3):
@@ -11,6 +12,38 @@
 fig = {"layout": {"title": {"text": "figure title"}}}
 
 
+def make_writeable_mocks():
+
+    # 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
+
+    # 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,16 +77,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,
+        )
 
-    count = writeable_mock.write.call_count + writeable_mock.write_bytes.call_count
-    assert count == 1
+        assert writeable_mock.active_write_function.call_count == 1
 
 
 def test_kaleido_engine_to_image_kwargs():
@@ -74,25 +110,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
-    )
-
-    count = writeable_mock.write.call_count + writeable_mock.write_bytes.call_count
-    assert count == 1
+        assert writeable_mock.active_write_function.call_count == 1
 
 
 def test_image_renderer():