Handle pathlib.Path in pio (#2974)

* Handle pathlib.Path in write_image

* Allow pathlib.Path in write_html

* Make file URI from absolute path (bugfix)

* pathlib for orca

* pathlib.Path support for read_json and write_json

Author: maresb

packages/python/plotly/plotly/basedatatypes.py

@@ -3578,7 +3578,7 @@ def write_html(self, *args, **kwargs):
         ----------
         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)
         config: dict or None (default None)
             Plotly.js figure config options
         auto_play: bool (default=True)
@@ -3751,7 +3751,7 @@ def write_image(self, *args, **kwargs):
         ----------
         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

packages/python/plotly/plotly/io/_html.py

@@ -1,6 +1,7 @@
 import uuid
 import json
 import os
+from pathlib import Path
 import webbrowser
 
 import six
@@ -401,7 +402,7 @@ def write_html(
         Figure object or dict representing a figure
     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)
     config: dict or None (default None)
         Plotly.js figure config options
     auto_play: bool (default=True)
@@ -520,24 +521,31 @@ def write_html(
     )
 
     # Check if file is a string
-    file_is_str = isinstance(file, six.string_types)
+    if isinstance(file, six.string_types):
+        # Use the standard pathlib constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, Path):  # 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
 
     # Write HTML string
-    if file_is_str:
-        with open(file, "w") as f:
-            f.write(html_str)
+    if path is not None:
+        path.write_text(html_str)
     else:
         file.write(html_str)
 
     # Check if we should copy plotly.min.js to output directory
-    if file_is_str and full_html and include_plotlyjs == "directory":
-        bundle_path = os.path.join(os.path.dirname(file), "plotly.min.js")
+    if path is not None and full_html and include_plotlyjs == "directory":
+        bundle_path = path.parent / "plotly.min.js"
 
-        if not os.path.exists(bundle_path):
-            with open(bundle_path, "w") as f:
-                f.write(get_plotlyjs())
+        if not bundle_path.exists():
+            bundle_path.write_text(get_plotlyjs())
 
     # Handle auto_open
-    if file_is_str and full_html and auto_open:
-        url = "file://" + os.path.abspath(file)
+    if path is not None and full_html and auto_open:
+        url = path.absolute().as_uri()
         webbrowser.open(url)

packages/python/plotly/plotly/io/_json.py

@@ -2,7 +2,7 @@
 
 from six import string_types
 import json
-
+from pathlib import Path
 
 from plotly.io._utils import validate_coerce_fig_to_dict, validate_coerce_output_type
 
@@ -68,7 +68,7 @@ def write_json(fig, file, validate=True, pretty=False, remove_uids=True):
 
     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)
 
     pretty: bool (default False)
         True if JSON representation should be pretty-printed, False if
@@ -87,17 +87,40 @@ def write_json(fig, file, validate=True, pretty=False, remove_uids=True):
     # Pass through validate argument and let to_json handle validation logic
     json_str = to_json(fig, validate=validate, pretty=pretty, remove_uids=remove_uids)
 
-    # 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 Path constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, Path):
+        # `file` is already a Path object.
+        path = file
+    else:
+        # We could not make a Path object out of file. Either `file` is an open file
+        # descriptor with a `write()` method or it's an invalid object.
+        path = None
 
     # Open file
     # ---------
-    if file_is_str:
-        with open(file, "w") as f:
-            f.write(json_str)
+    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(json_str)
+            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(json_str)
+        # We previously succeeded in interpreting `file` as a pathlib object.
+        # Now we can use `write_bytes()`.
+        path.write_text(json_str)
 
 
 def from_json(value, output_type="Figure", skip_invalid=False):
@@ -162,7 +185,7 @@ def read_json(file, output_type="Figure", skip_invalid=False):
     ----------
     file: str or readable
        A string containing the path to a local file or a read-able Python
-       object (e.g. an open file descriptor)
+       object (e.g. a pathlib.Path object or an open file descriptor)
 
     output_type: type or str (default 'Figure')
         The output figure type or type name.
@@ -177,17 +200,25 @@ def read_json(file, output_type="Figure", skip_invalid=False):
     Figure or FigureWidget
     """
 
-    # Check if file is a string
+    # Try to cast `file` as a pathlib object `path`.
     # -------------------------
-    # If it's a string we assume it's a local file path. If it's not a string
-    # then we assume it's a read-able Python object
+    # ----------------------------------------------
     file_is_str = isinstance(file, string_types)
+    if isinstance(file, string_types):
+        # Use the standard Path constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, Path):
+        # `file` is already a Path object.
+        path = file
+    else:
+        # We could not make a Path object out of file. Either `file` is an open file
+        # descriptor with a `write()` method or it's an invalid object.
+        path = None
 
     # Read file contents into JSON string
     # -----------------------------------
-    if file_is_str:
-        with open(file, "r") as f:
-            json_str = f.read()
+    if path is not None:
+        json_str = path.read_text()
     else:
         json_str = file.read()
 

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

@@ -2,6 +2,7 @@
 from six import string_types
 import os
 import json
+from pathlib import Path
 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 Path constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, Path):
+        # `file` is already a Path object.
+        path = file
+    else:
+        # We could not make a Path 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):

packages/python/plotly/plotly/io/_orca.py

@@ -10,6 +10,7 @@
 import warnings
 from copy import copy
 from contextlib import contextmanager
+from pathlib import Path
 
 import tenacity
 from six import string_types
@@ -1695,7 +1696,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
@@ -1741,16 +1742,25 @@ 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 Path constructor to make a pathlib object.
+        path = Path(file)
+    elif isinstance(file, Path):
+        # `file` is already a Path object.
+        path = file
+    else:
+        # We could not make a Path 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 = validate_coerce_format(ext)
+            format = ext.lstrip(".")
         else:
             raise ValueError(
                 """
@@ -1774,8 +1784,22 @@ 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)

packages/python/plotly/plotly/tests/test_io/test_pathlib.py

@@ -0,0 +1,57 @@
+"""Test compatibility with pathlib.Path.
+
+See also relevant tests in
+   packages/python/plotly/plotly/tests/test_optional/test_kaleido/test_kaleido.py
+"""
+
+from unittest import mock
+import plotly.io as pio
+from io import StringIO
+from pathlib import Path
+import re
+from unittest.mock import Mock
+
+fig = {"layout": {"title": {"text": "figure title"}}}
+
+
+def test_write_html():
+    """Verify that various methods for producing HTML have equivalent results.
+
+    The results will not be identical because the div id is pseudorandom. Thus
+    we compare the results after replacing the div id.
+
+    We test the results of
+    - pio.to_html
+    - pio.write_html with a StringIO buffer
+    - pio.write_html with a mock pathlib Path
+    - pio.write_html with a mock file descriptor
+    """
+    # Test pio.to_html
+    html = pio.to_html(fig)
+
+    # Test pio.write_html with a StringIO buffer
+    sio = StringIO()
+    pio.write_html(fig, sio)
+    sio.seek(0)  # Rewind to the beginning of the buffer, otherwise read() returns ''.
+    sio_html = sio.read()
+    assert replace_div_id(html) == replace_div_id(sio_html)
+
+    # Test pio.write_html with a mock pathlib Path
+    mock_pathlib_path = Mock(spec=Path)
+    pio.write_html(fig, mock_pathlib_path)
+    mock_pathlib_path.write_text.assert_called_once()
+    (pl_html,) = mock_pathlib_path.write_text.call_args[0]
+    assert replace_div_id(html) == replace_div_id(pl_html)
+
+    # Test pio.write_html with a mock file descriptor
+    mock_file_descriptor = Mock()
+    del mock_file_descriptor.write_bytes
+    pio.write_html(fig, mock_file_descriptor)
+    mock_file_descriptor.write.assert_called_once()
+    (fd_html,) = mock_file_descriptor.write.call_args[0]
+    assert replace_div_id(html) == replace_div_id(fd_html)
+
+
+def replace_div_id(s):
+    uuid = re.search(r'<div id="([^"]*)"', s).groups()[0]
+    return s.replace(uuid, "XXXX")