Merge pull request #43 from jakevdp/available-formats

ENH: add available_formats() function
altair-viz · Mar 21, 2020 · a5d9c7b · a5d9c7b
2 parents cef1b1c + ddadd38
commit a5d9c7b
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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("<!DOCTYPE html>")
+
+
+@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