Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit e69c31e

Browse files
committed
Handle pathlib.Path in write_image
1 parent b343fcc commit e69c31e

File tree

3 files changed

+142
-40
lines changed

3 files changed

+142
-40
lines changed

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

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from six import string_types
33
import os
44
import json
5+
from pathlib import Path, PurePath
56
import plotly
67
from plotly.io._utils import validate_coerce_fig_to_dict
78

@@ -169,7 +170,7 @@ def write_image(
169170
170171
file: str or writeable
171172
A string representing a local file path or a writeable object
172-
(e.g. an open file descriptor)
173+
(e.g. a pathlib.Path object or an open file descriptor)
173174
174175
format: str or None
175176
The desired image format. One of
@@ -228,14 +229,23 @@ def write_image(
228229
-------
229230
None
230231
"""
231-
# Check if file is a string
232-
# -------------------------
233-
file_is_str = isinstance(file, string_types)
232+
# Try to cast `file` as a pathlib object `path`.
233+
# ----------------------------------------------
234+
if isinstance(file, string_types):
235+
# Use the standard pathlib constructor to make a pathlib object.
236+
path = Path(file)
237+
elif isinstance(file, PurePath): # PurePath is the most general pathlib object.
238+
# `file` is already a pathlib object.
239+
path = file
240+
else:
241+
# We could not make a pathlib object out of file. Either `file` is an open file
242+
# descriptor with a `write()` method or it's an invalid object.
243+
path = None
234244

235245
# Infer format if not specified
236246
# -----------------------------
237-
if file_is_str and format is None:
238-
_, ext = os.path.splitext(file)
247+
if pathisnotNone and format is None:
248+
ext = path.suffix
239249
if ext:
240250
format = ext.lstrip(".")
241251
else:
@@ -267,11 +277,25 @@ def write_image(
267277

268278
# Open file
269279
# ---------
270-
if file_is_str:
271-
with open(file, "wb") as f:
272-
f.write(img_data)
280+
if path is None:
281+
# We previously failed to make sense of `file` as a pathlib object.
282+
# Attempt to write to `file` as an open file descriptor.
283+
try:
284+
file.write(img_data)
285+
return
286+
except AttributeError:
287+
pass
288+
raise ValueError(
289+
"""
290+
The 'file' argument '{file}' is not a string, pathlib.Path object, or file descriptor.
291+
""".format(
292+
file=file
293+
)
294+
)
273295
else:
274-
file.write(img_data)
296+
# We previously succeeded in interpreting `file` as a pathlib object.
297+
# Now we can use `write_bytes()`.
298+
path.write_bytes(img_data)
275299

276300

277301
def full_figure_for_development(fig, warn=True, as_dict=False):

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

Lines changed: 87 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,56 @@
11
import plotly.io as pio
22
import plotly.io.kaleido
3-
import sys
43
from contextlib import contextmanager
5-
6-
if sys.version_info >= (3, 3):
7-
from unittest.mock import Mock
8-
else:
9-
from mock import Mock
4+
from io import BytesIO
5+
from pathlib import Path
6+
from unittest.mock import Mock
107

118
fig = {"layout": {"title": {"text": "figure title"}}}
129

1310

11+
def make_writeable_mocks():
12+
"""Produce some mocks which we will use for testing the `write_image()` function.
13+
14+
These mocks should be passed as the `file=` argument to `write_image()`.
15+
16+
The tests should verify that the method specified in the `active_write_function`
17+
attribute is called once, and that scope.transform is called with the `format=`
18+
argument specified by the `.expected_format` attribute.
19+
20+
In total we provide two mocks: one for a writable file descriptor, and other for a
21+
pathlib.Path object.
22+
"""
23+
24+
# Part 1: A mock for a file descriptor
25+
# ------------------------------------
26+
mock_file_descriptor = Mock()
27+
28+
# A file descriptor has no write_bytes method, unlike a pathlib Path.
29+
del mock_file_descriptor.write_bytes
30+
31+
# The expected write method for a file descriptor is .write
32+
mock_file_descriptor.active_write_function = mock_file_descriptor.write
33+
34+
# Since there is no filename, there should be no format detected.
35+
mock_file_descriptor.expected_format = None
36+
37+
# Part 2: A mock for a pathlib path
38+
# ---------------------------------
39+
mock_pathlib_path = Mock(spec=Path)
40+
41+
# A pathlib Path object has no write method, unlike a file descriptor.
42+
del mock_pathlib_path.write
43+
44+
# The expected write method for a pathlib Path is .write_bytes
45+
mock_pathlib_path.active_write_function = mock_pathlib_path.write_bytes
46+
47+
# Mock a path with PNG suffix
48+
mock_pathlib_path.suffix = ".png"
49+
mock_pathlib_path.expected_format = "png"
50+
51+
return mock_file_descriptor, mock_pathlib_path
52+
53+
1454
@contextmanager
1555
def mocked_scope():
1656
# Code to acquire resource, e.g.:
@@ -44,15 +84,19 @@ def test_kaleido_engine_to_image():
4484

4585

4686
def test_kaleido_engine_write_image():
47-
writeable_mock =Mock()
48-
with mocked_scope() as scope:
49-
pio.write_image(fig, writeable_mock, engine="kaleido", validate=False)
87+
forwriteable_mock inmake_writeable_mocks():
88+
with mocked_scope() as scope:
89+
pio.write_image(fig, writeable_mock, engine="kaleido", validate=False)
5090

51-
scope.transform.assert_called_with(
52-
fig, format=None, width=None, height=None, scale=None
53-
)
91+
scope.transform.assert_called_with(
92+
fig,
93+
format=writeable_mock.expected_format,
94+
width=None,
95+
height=None,
96+
scale=None,
97+
)
5498

55-
assert writeable_mock.write.call_count == 1
99+
assert writeable_mock.active_write_function.call_count == 1
56100

57101

58102
def test_kaleido_engine_to_image_kwargs():
@@ -73,24 +117,24 @@ def test_kaleido_engine_to_image_kwargs():
73117

74118

75119
def test_kaleido_engine_write_image_kwargs():
76-
writeable_mock = Mock()
77-
with mocked_scope() as scope:
78-
pio.write_image(
79-
fig,
80-
writeable_mock,
81-
format="jpg",
82-
width=700,
83-
height=600,
84-
scale=2,
85-
engine="kaleido",
86-
validate=False,
120+
for writeable_mock in make_writeable_mocks():
121+
with mocked_scope() as scope:
122+
pio.write_image(
123+
fig,
124+
writeable_mock,
125+
format="jpg",
126+
width=700,
127+
height=600,
128+
scale=2,
129+
engine="kaleido",
130+
validate=False,
131+
)
132+
133+
scope.transform.assert_called_with(
134+
fig, format="jpg", width=700, height=600, scale=2
87135
)
88136

89-
scope.transform.assert_called_with(
90-
fig, format="jpg", width=700, height=600, scale=2
91-
)
92-
93-
assert writeable_mock.write.call_count == 1
137+
assert writeable_mock.active_write_function.call_count == 1
94138

95139

96140
def test_image_renderer():
@@ -105,3 +149,17 @@ def test_image_renderer():
105149
height=renderer.height,
106150
scale=renderer.scale,
107151
)
152+
153+
154+
def test_bytesio():
155+
"""Verify that writing to a BytesIO object contains the same data as to_image().
156+
157+
The goal of this test is to ensure that Plotly correctly handles a writable buffer
158+
which doesn't correspond to a filesystem path.
159+
"""
160+
bio = BytesIO()
161+
pio.write_image(fig, bio, format="jpg", engine="kaleido", validate=False)
162+
bio.seek(0) # Rewind to the beginning of the buffer, otherwise read() returns b''.
163+
bio_bytes = bio.read()
164+
to_image_bytes = pio.to_image(fig, format="jpg", engine="kaleido", validate=False)
165+
assert bio_bytes == to_image_bytes

‎packages/python/plotly/plotly/tests/test_orca/test_to_image.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import pytest
77
import sys
88
import pandas as pd
9+
from io import BytesIO
910

1011
if sys.version_info >= (3, 3):
1112
from unittest.mock import MagicMock
@@ -234,7 +235,12 @@ def test_write_image_writeable(fig1, format):
234235
fig1, mock_file, format=format, width=700, height=500, engine="orca"
235236
)
236237

237-
mock_file.write.assert_called_once_with(expected_bytes)
238+
if mock_file.write_bytes.called:
239+
mock_file.write_bytes.assert_called_once_with(expected_bytes)
240+
elif mock_file.write.called:
241+
mock_file.write.assert_called_once_with(expected_bytes)
242+
else:
243+
assert "Neither write nor write_bytes was called."
238244

239245

240246
def test_write_image_string_format_inference(fig1, format):
@@ -347,3 +353,17 @@ def test_invalid_figure_json():
347353
)
348354

349355
assert "400: invalid or malformed request syntax" in str(err.value)
356+
357+
358+
def test_bytesio(fig1):
359+
"""Verify that writing to a BytesIO object contains the same data as to_image().
360+
361+
The goal of this test is to ensure that Plotly correctly handles a writable buffer
362+
which doesn't correspond to a filesystem path.
363+
"""
364+
bio = BytesIO()
365+
pio.write_image(fig1, bio, format="jpg", validate=False)
366+
bio.seek(0)
367+
bio_bytes = bio.read()
368+
to_image_bytes = pio.to_image(fig1, format="jpg", validate=False)
369+
assert bio_bytes == to_image_bytes

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /