def test_extract_format(ext: str, fmt: str, use_filename: bool) -> None:
if use_filename:
filename = f"chart.{ext}"
assert extract_format(filename) == fmt
else:
with tempfile.NamedTemporaryFile(suffix=f".{ext}") as fp:
assert extract_format(fp) == fmt
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 _get_saver_for_format(
fmt: Optional[str] = None, fp: Optional[Union[IO, str]] = None
) -> Type[Saver]:
"""Get an enabled Saver class that supports the specified format."""
# TODO: allow other savers to be registered.
if fmt is None:
if fp is None:
raise ValueError("Either fmt or fp must be specified")
fmt = extract_format(fp)
savers: List[Type[Saver]] = [BasicSaver, HTMLSaver, SeleniumSaver, NodeSaver]
for s in savers:
if fmt in s.valid_formats and s.enabled():
return s
raise ValueError(f"Unsupported format: {fmt!r}")
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 _select_saver(
method: Optional[Union[str, Type[Saver]]],
mode: str,
fmt: Optional[str] = None,
fp: Optional[Union[IO, str]] = None,
) -> Type[Saver]:
"""Get an enabled Saver class that supports the specified format.
Parameters
----------
method : string or Saver class or None
The saver class to use. If None, the saver class will be chosen
automatically.
mode : string
One of "vega" or "vega-lite".
fmt : string, optional
The format to which the spec will be saved. If not specified, it
is inferred from `fp`.
fp : string or file-like object, optional
Only referenced if fmt is None. The name is used to infer the format
if possible.
Returns
-------
Saver : Saver class
The Saver subclass that implements the desired operation.
"""
if isinstance(method, type) and issubclass(method, Saver):
return method
elif isinstance(method, str):
if method in _SAVER_METHODS:
return _SAVER_METHODS[method]
else:
raise ValueError(f"Unrecognized method: {method!r}")
elif method is None:
if fmt is None:
if fp is None:
raise ValueError("Either fmt or fp must be specified")
fmt = extract_format(fp)
for s in _SAVER_METHODS.values():
if s.enabled() and fmt in s.valid_formats[mode]:
return s
raise ValueError(
f"No enabled saver found that supports format={fmt!r}")
else:
raise ValueError(f"Unrecognized method: {method}")
def test_extract_format_failure() -> None:
fp = io.StringIO()
with pytest.raises(ValueError) as err:
extract_format(fp)
assert f"Cannot infer format from {fp}" in str(err.value)