def save(self, fp: Union[IO, str], fmt: Optional[str] = None) -> None:
"""Save a chart to file
Parameters
----------
fp : file or filename
Location to save the result. For fmt in ["png", "pdf"], file must be binary.
For fmt in ["svg", "vega", "vega-lite"], file must be text.
fmt : string
The format in which to save the chart. If not specified and fp is a string,
fmt will be determined from the file extension.
"""
if fmt is None:
fmt = extract_format(fp)
if fmt not in self.valid_formats:
raise ValueError(
f"Got fmt={fmt}; expected one of {self.valid_formats}")
content = self.mimebundle(fmt).popitem()[1]
if isinstance(content, dict):
with maybe_open(fp, "w") as f:
json.dump(content, f, indent=2)
elif isinstance(content, str):
with maybe_open(fp, "w") as f:
f.write(content)
elif isinstance(content, bytes):
with maybe_open(fp, "wb") as f:
f.write(content)
else:
raise ValueError(
f"Unrecognized content type: {type(content)} for fmt={fmt!r}")
def test_maybe_open_errors() -> None:
with pytest.raises(ValueError) as err:
with maybe_open(io.BytesIO(), "w"):
pass
assert "fp is opened in binary mode" in str(err.value)
assert "mode='w'" in str(err.value)
with pytest.raises(ValueError) as err:
with maybe_open(io.StringIO(), "wb"):
pass
assert "fp is opened in text mode" in str(err.value)
assert "mode='wb'" in str(err.value)
def save(
self, fp: Optional[Union[IO, str]] = None, fmt: Optional[str] = None
) -> Optional[Union[str, bytes]]:
"""Save a chart to file
Parameters
----------
fp : file or filename (optional)
Location to save the result. For fmt in ["png", "pdf"], file must be binary.
For fmt in ["svg", "vega", "vega-lite"], file must be text. If not specified,
the serialized chart will be returned.
fmt : string (optional)
The format in which to save the chart. If not specified and fp is a string,
fmt will be determined from the file extension.
Returns
-------
chart : string, bytes, or None
If fp is None, the serialized chart is returned.
If fp is specified, the return value is None.
"""
if fmt is None:
if fp is None:
raise ValueError("Must specify either `fp` or `fmt` when saving chart")
fmt = extract_format(fp)
if fmt not in self.valid_formats[self._mode]:
raise ValueError(f"Got fmt={fmt}; expected one of {self.valid_formats}")
content = self._serialize(fmt, "save")
if fp is None:
if isinstance(content, dict):
return json.dumps(content)
return content
if isinstance(content, dict):
with maybe_open(fp, "w") as f:
json.dump(content, f, indent=2)
elif isinstance(content, str):
with maybe_open(fp, "w") as f:
f.write(content)
elif isinstance(content, bytes):
with maybe_open(fp, "wb") as f:
f.write(content)
else:
raise ValueError(
f"Unrecognized content type: {type(content)} for fmt={fmt!r}"
)
return None
def test_maybe_open_fileobj(mode: str) -> None:
content_raw = "testing maybe_open with file object\n"
content = content_raw.encode() if "b" in mode else content_raw
with tempfile.NamedTemporaryFile(mode + "+") as fp:
with maybe_open(fp, mode) as f:
f.write(content)
fp.seek(0)
assert fp.read() == content
def test_maybe_open_filename(mode: str) -> None:
content_raw = "testing maybe_open with filename\n"
content = content_raw.encode() if "b" in mode else content_raw
with temporary_filename() as filename:
with maybe_open(filename, mode) as f:
f.write(content)
with open(filename, "rb" if "b" in mode else "r") as f:
assert f.read() == content