File size: 4,455 Bytes
ba2f5d6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import json
import pathlib

from .mimebundle import spec_to_mimebundle


def write_file_or_filename(fp, content, mode="w"):
    """Write content to fp, whether fp is a string, a pathlib Path or a
    file-like object"""
    if isinstance(fp, str) or isinstance(fp, pathlib.PurePath):
        with open(fp, mode) as f:
            f.write(content)
    else:
        fp.write(content)


def save(
    chart,
    fp,
    vega_version,
    vegaembed_version,
    format=None,
    mode=None,
    vegalite_version=None,
    embed_options=None,
    json_kwds=None,
    webdriver="chrome",
    scale_factor=1,
    **kwargs,
):
    """Save a chart to file in a variety of formats

    Supported formats are [json, html, png, svg]

    Parameters
    ----------
    chart : alt.Chart
        the chart instance to save
    fp : string filename, pathlib.Path or file-like object
        file to which to write the chart.
    format : string (optional)
        the format to write: one of ['json', 'html', 'png', 'svg'].
        If not specified, the format will be determined from the filename.
    mode : string (optional)
        Either 'vega' or 'vegalite'. If not specified, then infer the mode from
        the '$schema' property of the spec, or the ``opt`` dictionary.
        If it's not specified in either of those places, then use 'vegalite'.
    vega_version : string
        For html output, the version of vega.js to use
    vegalite_version : string
        For html output, the version of vegalite.js to use
    vegaembed_version : string
        For html output, the version of vegaembed.js to use
    embed_options : dict
        The vegaEmbed options dictionary. Default is {}
        (See https://github.com/vega/vega-embed for details)
    json_kwds : dict
        Additional keyword arguments are passed to the output method
        associated with the specified format.
    webdriver : string {'chrome' | 'firefox'}
        Webdriver to use for png or svg output
    scale_factor : float
        scale_factor to use to change size/resolution of png or svg output
    **kwargs :
        additional kwargs passed to spec_to_mimebundle.
    """
    if json_kwds is None:
        json_kwds = {}

    if embed_options is None:
        embed_options = {}

    if format is None:
        if isinstance(fp, str):
            format = fp.split(".")[-1]
        elif isinstance(fp, pathlib.PurePath):
            format = fp.suffix.lstrip(".")
        else:
            raise ValueError(
                "must specify file format: " "['png', 'svg', 'pdf', 'html', 'json']"
            )

    spec = chart.to_dict()

    if mode is None:
        if "mode" in embed_options:
            mode = embed_options["mode"]
        elif "$schema" in spec:
            mode = spec["$schema"].split("/")[-2]
        else:
            mode = "vega-lite"

    if mode not in ["vega", "vega-lite"]:
        raise ValueError("mode must be 'vega' or 'vega-lite', " "not '{}'".format(mode))

    if mode == "vega-lite" and vegalite_version is None:
        raise ValueError("must specify vega-lite version")

    if format == "json":
        json_spec = json.dumps(spec, **json_kwds)
        write_file_or_filename(fp, json_spec, mode="w")
    elif format == "html":
        mimebundle = spec_to_mimebundle(
            spec=spec,
            format=format,
            mode=mode,
            vega_version=vega_version,
            vegalite_version=vegalite_version,
            vegaembed_version=vegaembed_version,
            embed_options=embed_options,
            json_kwds=json_kwds,
            **kwargs,
        )
        write_file_or_filename(fp, mimebundle["text/html"], mode="w")
    elif format in ["png", "svg", "pdf"]:
        mimebundle = spec_to_mimebundle(
            spec=spec,
            format=format,
            mode=mode,
            vega_version=vega_version,
            vegalite_version=vegalite_version,
            vegaembed_version=vegaembed_version,
            webdriver=webdriver,
            scale_factor=scale_factor,
            **kwargs,
        )
        if format == "png":
            write_file_or_filename(fp, mimebundle["image/png"], mode="wb")
        elif format == "pdf":
            write_file_or_filename(fp, mimebundle["application/pdf"], mode="wb")
        else:
            write_file_or_filename(fp, mimebundle["image/svg+xml"], mode="w")
    else:
        raise ValueError("unrecognized format: '{}'".format(format))