diff --git a/CHANGES.md b/CHANGES.md index 0c18e09..21009fb 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -2,6 +2,9 @@ ## Version 0.4.0 (unreleased) +- Added top-level ``available_formats()`` function, which returns the set of + available formats. (#43) + ## Version 0.3.1 - Fix bug in detecting npm binary path (#42) diff --git a/altair_saver/__init__.py b/altair_saver/__init__.py index de10b8e..684294d 100644 --- a/altair_saver/__init__.py +++ b/altair_saver/__init__.py @@ -1,9 +1,10 @@ """Tools for saving altair charts""" -from ._core import render, save +from ._core import render, save, available_formats from .savers import Saver, BasicSaver, HTMLSaver, NodeSaver, SeleniumSaver __version__ = "0.4.0.dev0" __all__ = [ + "available_formats", "render", "save", "Saver", diff --git a/altair_saver/_core.py b/altair_saver/_core.py index d4e3c18..0e34950 100644 --- a/altair_saver/_core.py +++ b/altair_saver/_core.py @@ -1,5 +1,5 @@ from collections import OrderedDict -from typing import Any, Dict, IO, Iterable, Optional, Type, Union +from typing import Any, Dict, IO, Iterable, Optional, Set, Type, Union import altair as alt @@ -239,3 +239,24 @@ def render( mimebundle.update(saver.mimebundle(fmt)) return mimebundle + + +def available_formats(mode: str = "vega-lite") -> Set[str]: + """Return the set of available formats. + + Parameters + ---------- + mode : str + The kind of input; one of "vega", "vega-lite" + + Returns + ------- + formats : set of strings + Formats available in the current session. + """ + valid_modes = {"vega", "vega-lite"} + if mode not in valid_modes: + raise ValueError(f"Invalid mode: {mode!r}. Must be one of {valid_modes!r}") + return set.union( + *(set(s.valid_formats[mode]) for s in _SAVER_METHODS.values() if s.enabled()) + ) diff --git a/altair_saver/tests/test_core.py b/altair_saver/tests/test_core.py index 0f69427..0873ff4 100644 --- a/altair_saver/tests/test_core.py +++ b/altair_saver/tests/test_core.py @@ -1,12 +1,13 @@ import io import json -from typing import Dict, List, Union, Type +from typing import Any, Dict, List, Union, Type import altair as alt import pandas as pd import pytest from altair_saver import ( + available_formats, save, render, BasicSaver, @@ -186,3 +187,19 @@ def test_infer_format(spec: JSONDict) -> None: with open(filename, "r") as fp: html = fp.read() assert html.strip().startswith("") + + +@pytest.mark.parametrize("mode", ["vega", "vega-lite"]) +def test_available_formats(monkeypatch: Any, mode: str) -> None: + monkeypatch.setattr(NodeSaver, "enabled", lambda: False) + monkeypatch.setattr(SeleniumSaver, "enabled", lambda: False) + expected = {mode, "json", "html"} + assert available_formats(mode) == expected + + monkeypatch.setattr(SeleniumSaver, "enabled", lambda: True) + expected |= {"vega", "png", "svg"} + assert available_formats(mode) == expected + + monkeypatch.setattr(NodeSaver, "enabled", lambda: True) + expected |= {"pdf"} + assert available_formats(mode) == expected