Repository: altair-viz/altair
Pull request #556 (2018-03-06 04:55:54)

Add savechart() method to vegalite/v2


Patch:
diff --git a/altair/utils/core.py b/altair/utils/core.py
index c3398759..8bf7b056 100644
--- a/altair/utils/core.py
+++ b/altair/utils/core.py
@@ -339,3 +339,12 @@ def update_nested(original, update, copy=False):
         else:
             original[key] = val
     return original
+
+
+def write_file_or_filename(fp, content, mode='w'):
+    """Write content to fp, whether fp is a string or a file-like object"""
+    if isinstance(fp, six.string_types):
+        with open(fp, mode) as f:
+            f.write(content)
+    else:
+        fp.write(content)
diff --git a/altair/utils/headless.py b/altair/utils/headless.py
index 9aeba70e..28769c7b 100644
--- a/altair/utils/headless.py
+++ b/altair/utils/headless.py
@@ -8,7 +8,10 @@ import json
 import os
 import tempfile
 
+import six
+
 from .importing import attempt_import
+from .core import write_file_or_filename
 
 
 HTML_TEMPLATE = """
@@ -27,55 +30,71 @@ HTML_TEMPLATE = """
 """
 
 EMBED_CODE = """
-vegaEmbed("#vis", {spec}).then(function(result) {{
+const spec = {spec};
+const opt = {opt};
+vegaEmbed("#vis", spec, opt).then(function(result) {{
     window.view = result.view;
 }}).catch(console.error);
 """
 
-CONVERT_CODE = """
-window.view.toCanvas().then(function(canvas) {
-    window.png_render = canvas.toDataURL('image/png');
-})
-"""
+CONVERT_CODE = {
+'png': """
+       window.view.toCanvas().then(function(canvas) {
+           window.image_render = canvas.toDataURL('image/png');
+       })
+       """,
+'svg': """
+       window.view.toSVG().then(function(render) {
+           window.image_render = render;
+       })
+       """}
 
 EXTRACT_CODE = """
-return window.png_render;
+return window.image_render;
 """
 
-def save_spec(spec, filename, format=None, mode='vega-lite'):
+def save_spec(spec, fp, mode=None, format=None):
     """Save a spec to file
 
     Parameters
     ----------
     spec : dict
         a dictionary representing a vega-lite plot spec
-    filename : string
-        the filename at which the result will be saved
+    fp : string or file-like object
+        the filename or file object at which the result will be saved
+    mode : string or None
+        The rendering mode ('vega' or 'vega-lite'). If None, the mode will be
+        inferred from the $schema attribute of the spec, or will default to
+        'vega' if $schema is not in the spec.
     format : string (optional)
         the file format to be saved. If not specified, it will be inferred
-        from the extension of filename
-    mode : string
-        Whether the spec is 'vega' or 'vega-lite'.
-        Currently only mode='vega-lite' is supported.
+        from the extension of filename.
 
     Note
     ----
     This requires the pillow, selenium, and chrome headless packages to be
     installed.
     """
-    # TODO: remove PIL dependency?
-    # TODO: use SVG renderer when it makes sense
-    # TODO: support mode='vega'
     # TODO: allow package versions to be specified
-    # TODO: detect local Jupyter caches of JS packages?
+    # TODO: detect & use local Jupyter caches of JS packages?
+
+    if format is None and isinstance(fp, six.string_types):
+        format = fp.split('.')[-1]
+
+    if format not in ['png', 'svg']:
+        raise NotImplementedError("save_spec only supports 'svg' and 'png'")
 
-    Image = attempt_import('PIL.Image',
-                           'save_spec requires the pillow package')
     webdriver = attempt_import('selenium.webdriver',
                                'save_spec requires the selenium package')
     Options = attempt_import('selenium.webdriver.chrome.options',
                              'save_spec requires the selenium package').Options
 
+    opt = {'renderer': 'canvas' if format == 'png' else 'svg'}
+    if mode is not None:
+        if mode not in ['vega', 'vega-lite']:
+            raise ValueError("mode must be 'vega' or 'vega-lite'")
+        opt['mode'] = mode
+
     try:
         chrome_options = Options()
         chrome_options.add_argument("--headless")
@@ -85,17 +104,18 @@ def save_spec(spec, filename, format=None, mode='vega-lite'):
             with open(name, 'w') as f:
                 f.write(HTML_TEMPLATE)
             driver.get("file://" + name)
-            driver.execute_script(EMBED_CODE.format(spec=json.dumps(spec)))
-            driver.execute_script(CONVERT_CODE)
-            png_base64 = driver.execute_script(EXTRACT_CODE)
+            driver.execute_script(EMBED_CODE.format(spec=json.dumps(spec),
+                                                    opt=json.dumps(opt)))
+            driver.execute_script(CONVERT_CODE[format])
+            render = driver.execute_script(EXTRACT_CODE)
         finally:
             os.remove(name)
     finally:
         driver.close()
 
-    out = io.BytesIO()
-    metadata, image = png_base64.split(',')
-    base64.decode(io.BytesIO(image.encode()), out)
-    out.seek(0)
-    image = Image.open(out)
-    image.save(filename)
+    if format == 'png':
+        img_bytes = base64.decodebytes(render.split(',')[1].encode())
+        write_file_or_filename(fp, img_bytes, mode='wb')
+    else:
+        img_bytes = render.encode()
+        write_file_or_filename(fp, img_bytes, mode='wb')
diff --git a/altair/utils/schemapi.py b/altair/utils/schemapi.py
index 32bac966..875bf665 100644
--- a/altair/utils/schemapi.py
+++ b/altair/utils/schemapi.py
@@ -3,7 +3,6 @@
 import collections
 import contextlib
 import json
-import textwrap
 
 import jsonschema
 import six
@@ -105,7 +104,7 @@ class SchemaBase(object):
         object.__setattr__(self, '_kwds', kwds)
 
         if DEBUG_MODE and self._class_is_valid_at_instantiation:
-            dct = self.to_dict(validate=True)
+            _ = self.to_dict(validate=True)
 
     def copy(self, deep=True, ignore=()):
         """Return a copy of the object
@@ -227,6 +226,38 @@ class SchemaBase(object):
                 raise SchemaValidationError(self, err)
         return result
 
+    def to_json(self, validate=True, ignore=[], context={},
+                indent=2, sort_keys=True, **kwargs):
+        """Emit the JSON representation for this object as a string.
+
+        Parameters
+        ----------
+        validate : boolean or string
+            If True (default), then validate the output dictionary
+            against the schema. If "deep" then recursively validate
+            all objects in the spec. This takes much more time, but
+            it results in friendlier tracebacks for large objects.
+        ignore : list
+            A list of keys to ignore. This will *not* passed to child to_dict
+            function calls.
+        context : dict (optional)
+            A context dictionary that will be passed to all child to_dict
+            function calls
+        indent : integer, default 2
+            the number of spaces of indentation to use
+        sort_keys : boolean, default True
+            if True, sort keys in the output
+        **kwargs
+            Additional keyword arguments are passed to ``json.dumps()``
+
+        Returns
+        -------
+        spec : string
+            The JSON specification of the chart object.
+        """
+        dct = self.to_dict(validate=validate, ignore=ignore, context=context)
+        return json.dumps(dct, indent=indent, sort_keys=sort_keys, **kwargs)
+
     @classmethod
     def from_dict(cls, dct, validate=True):
         """Construct class from a dictionary representation
@@ -253,6 +284,26 @@ class SchemaBase(object):
         converter = _FromDict(SchemaBase.__subclasses__())
         return converter.from_dict(constructor=cls, root=cls,
                                    schema=cls._schema, dct=dct)
+    @classmethod
+    def from_json(cls, json_string, validate=True, **kwargs):
+        """Instantiate the object from a valid JSON string
+
+        Parameters
+        ----------
+        json_string : string
+            The string containing a valid JSON chart specification.
+        validate : boolean
+            If True (default), then validate the input against the schema.
+        **kwargs :
+            Additional keyword arguments are passed to json.loads
+
+        Returns
+        -------
+        chart : Chart object
+            The altair Chart object built from the specification.
+        """
+        dct = json.loads(json_string, **kwargs)
+        return cls.from_dict(dct, validate=validate)
 
     @classmethod
     def validate(cls, instance, schema=None):
diff --git a/altair/vegalite/v2/api.py b/altair/vegalite/v2/api.py
index a5b7fb55..d2092f5d 100644
--- a/altair/vegalite/v2/api.py
+++ b/altair/vegalite/v2/api.py
@@ -1,3 +1,5 @@
+import json
+
 import jsonschema
 import six
 import pandas as pd
@@ -7,7 +9,8 @@ from .schema import core, channels, mixins, Undefined
 from .data import data_transformers, pipe
 from ...utils import (infer_vegalite_type, parse_shorthand,
                       parse_shorthand_plus_data,
-                      use_signature, update_nested)
+                      use_signature, update_nested,
+                      save_spec, write_file_or_filename)
 from .display import renderers
 
 
@@ -195,9 +198,10 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
         # - 'top_level' is a boolean flag that is assumed to be true; if it's
         #   true then a "$schema" arg is added to the dict.
         context = kwargs.get('context', {}).copy()
-        is_top_level = context.get('top_level', True)
 
+        is_top_level = context.get('top_level', True)
         context['top_level'] = False
+
         if original_data is not Undefined:
             context['data'] = original_data
         kwargs['context'] = context
@@ -208,7 +212,8 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
             dct = None
 
         # If we hit an error, then re-convert with validate='deep' to get
-        # a more useful traceback
+        # a more useful traceback. We don't do this by default because it's
+        # much slower in the case that there are no errors.
         if dct is None:
             kwargs['validate'] = 'deep'
             dct = super(TopLevelMixin, copy).to_dict(*args, **kwargs)
@@ -223,6 +228,45 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
                 dct = update_nested(copy._default_spec_values, dct, copy=True)
         return dct
 
+    def savechart(self, fp, format=None, **kwargs):
+        """Save a chart to file in a variety of formats
+
+        Supported formats are json, html, png, svg
+
+        Parameters
+        ----------
+        fp : string filename or file-like object
+            file in which to write the chart.
+        format : string (optional)
+            the format to write: one of ['json', 'html', 'png', 'eps'].
+            If not specified, the format will be determined from the filename.
+        **kwargs :
+            Additional keyword arguments are passed to the output method
+            associated with the specified format.
+        """
+
+        if isinstance(fp, six.string_types):
+            format = fp.split('.')[-1]
+
+        if format is None:
+            raise ValueError("must specify file format: "
+                             "['png', 'eps', 'html', 'json']")
+        elif format == 'json':
+            write_file_or_filename(fp, self.to_json(**kwargs), mode='w')
+        elif format == 'html':
+            from .html import HTML_TEMPLATE
+            opt = dict(renderer=kwargs.pop('renderer', 'canvas'),
+                       actions=kwargs.pop('actions', False))
+            if opt['renderer'] not in ('canvas', 'svg'):
+                raise ValueError("renderer must be 'canvas' or 'svg'")
+            spec_html = HTML_TEMPLATE.format(spec=self.to_json(**kwargs),
+                                             opt=json.dumps(opt))
+            write_file_or_filename(fp, spec_html, mode='w')
+        elif format in ['png', 'svg']:
+            save_spec(self.to_dict(), fp, format=format, **kwargs)
+        else:
+            raise ValueError("unrecognized format: '{0}'".format(format))
+
     # Layering and stacking
 
     def __add__(self, other):
diff --git a/altair/vegalite/v2/html.py b/altair/vegalite/v2/html.py
new file mode 100644
index 00000000..1debdf0a
--- /dev/null
+++ b/altair/vegalite/v2/html.py
@@ -0,0 +1,18 @@
+HTML_TEMPLATE = """
+<!DOCTYPE html>
+<html>
+<head>
+  <script src="https://cdn.jsdelivr.net/npm/vega@3"></script>
+  <script src="https://cdn.jsdelivr.net/npm/vega-lite@2"></script>
+  <script src="https://cdn.jsdelivr.net/npm/vega-embed@3"></script>
+</head>
+<body>
+  <div id="vis"></div>
+  <script type="text/javascript">
+    var spec = {spec};
+    var opt = {opt};
+    vegaEmbed("#vis", spec, opt);
+  </script>
+</body>
+</html>
+"""
diff --git a/tools/schemapi/schemapi.py b/tools/schemapi/schemapi.py
index c6143bb1..cf915c56 100644
--- a/tools/schemapi/schemapi.py
+++ b/tools/schemapi/schemapi.py
@@ -1,7 +1,6 @@
 import collections
 import contextlib
 import json
-import textwrap
 
 import jsonschema
 import six
@@ -103,7 +102,7 @@ class SchemaBase(object):
         object.__setattr__(self, '_kwds', kwds)
 
         if DEBUG_MODE and self._class_is_valid_at_instantiation:
-            dct = self.to_dict(validate=True)
+            _ = self.to_dict(validate=True)
 
     def copy(self, deep=True, ignore=()):
         """Return a copy of the object
@@ -225,6 +224,38 @@ class SchemaBase(object):
                 raise SchemaValidationError(self, err)
         return result
 
+    def to_json(self, validate=True, ignore=[], context={},
+                indent=2, sort_keys=True, **kwargs):
+        """Emit the JSON representation for this object as a string.
+
+        Parameters
+        ----------
+        validate : boolean or string
+            If True (default), then validate the output dictionary
+            against the schema. If "deep" then recursively validate
+            all objects in the spec. This takes much more time, but
+            it results in friendlier tracebacks for large objects.
+        ignore : list
+            A list of keys to ignore. This will *not* passed to child to_dict
+            function calls.
+        context : dict (optional)
+            A context dictionary that will be passed to all child to_dict
+            function calls
+        indent : integer, default 2
+            the number of spaces of indentation to use
+        sort_keys : boolean, default True
+            if True, sort keys in the output
+        **kwargs
+            Additional keyword arguments are passed to ``json.dumps()``
+
+        Returns
+        -------
+        spec : string
+            The JSON specification of the chart object.
+        """
+        dct = self.to_dict(validate=validate, ignore=ignore, context=context)
+        return json.dumps(dct, indent=indent, sort_keys=sort_keys, **kwargs)
+
     @classmethod
     def from_dict(cls, dct, validate=True):
         """Construct class from a dictionary representation
@@ -251,6 +282,26 @@ class SchemaBase(object):
         converter = _FromDict(SchemaBase.__subclasses__())
         return converter.from_dict(constructor=cls, root=cls,
                                    schema=cls._schema, dct=dct)
+    @classmethod
+    def from_json(cls, json_string, validate=True, **kwargs):
+        """Instantiate the object from a valid JSON string
+
+        Parameters
+        ----------
+        json_string : string
+            The string containing a valid JSON chart specification.
+        validate : boolean
+            If True (default), then validate the input against the schema.
+        **kwargs :
+            Additional keyword arguments are passed to json.loads
+
+        Returns
+        -------
+        chart : Chart object
+            The altair Chart object built from the specification.
+        """
+        dct = json.loads(json_string, **kwargs)
+        return cls.from_dict(dct, validate=validate)
 
     @classmethod
     def validate(cls, instance, schema=None):


Test patch:
diff --git a/altair/utils/tests/test_schemapi.py b/altair/utils/tests/test_schemapi.py
index d9ba9166..07ecfe55 100644
--- a/altair/utils/tests/test_schemapi.py
+++ b/altair/utils/tests/test_schemapi.py
@@ -1,5 +1,7 @@
 # The contents of this file are automatically written by
 # tools/generate_schema_wrapper.py. Do not modify directly.
+import json
+
 import jsonschema
 import pytest
 
@@ -240,3 +242,13 @@ def test_attribute_error():
         m.invalid_attribute
     assert str(err.value) == ("'MySchema' object has no attribute "
                               "'invalid_attribute'")
+
+
+def test_to_from_json():
+    dct = {'a': {'foo': 'bar'}, 'a2': {'foo': 42},
+           'b': ['a', 'b', 'c'], 'b2': [1, 2, 3], 'c': 42,
+           'd': ['x', 'y', 'z']}
+    json_str = MySchema.from_dict(dct).to_json()
+    new_dct = MySchema.from_json(json_str).to_dict()
+
+    assert new_dct == dct
diff --git a/altair/vegalite/v2/tests/test_api.py b/altair/vegalite/v2/tests/test_api.py
index 63559ed6..a00da648 100644
--- a/altair/vegalite/v2/tests/test_api.py
+++ b/altair/vegalite/v2/tests/test_api.py
@@ -1,5 +1,10 @@
 """Unit tests for altair API"""
 
+import io
+import json
+import os
+import tempfile
+
 import pytest
 import pandas as pd
 
@@ -114,3 +119,43 @@ def test_selection_to_dict():
         text=alt.condition(brush, alt.value('abc'), alt.Text('col2:N')),
         size=alt.condition(brush, alt.value(20), 'col2:N')
     ).to_dict()
+
+
+@pytest.mark.parametrize('format', ['html', 'json', 'png', 'svg'])
+def test_savechart(format):
+    from ..examples.bar import chart
+
+    if format in ['html', 'json']:
+        out = io.StringIO()
+        mode = 'r'
+    else:
+        out = io.BytesIO()
+        mode = 'rb'
+    fid, filename = tempfile.mkstemp(suffix='.' + format)
+
+    try:
+        # selenium may not be installed; skip the test if we get a selenium error.
+        try:
+            chart.savechart(out, format=format)
+            chart.savechart(filename)
+        except RuntimeError as err:
+            if format in ['png', 'svg'] and 'selenium' in str(err):
+                pytest.skip("selenium installation required for png/svg export")
+            else:
+                raise
+
+        out.seek(0)
+        with open(filename, mode) as f:
+            assert f.read() == out.read()
+    finally:
+        os.remove(filename)
+
+    out.seek(0)
+
+    if format == 'json':
+        spec = json.load(out)
+        assert '$schema' in spec
+
+    elif format == 'html':
+        content = out.read()
+        assert content.startswith('\n<!DOCTYPE html>')
diff --git a/tools/schemapi/tests/test_schemapi.py b/tools/schemapi/tests/test_schemapi.py
index f231b002..534525a2 100644
--- a/tools/schemapi/tests/test_schemapi.py
+++ b/tools/schemapi/tests/test_schemapi.py
@@ -1,3 +1,5 @@
+import json
+
 import jsonschema
 import pytest
 
@@ -238,3 +240,13 @@ def test_attribute_error():
         m.invalid_attribute
     assert str(err.value) == ("'MySchema' object has no attribute "
                               "'invalid_attribute'")
+
+
+def test_to_from_json():
+    dct = {'a': {'foo': 'bar'}, 'a2': {'foo': 42},
+           'b': ['a', 'b', 'c'], 'b2': [1, 2, 3], 'c': 42,
+           'd': ['x', 'y', 'z']}
+    json_str = MySchema.from_dict(dct).to_json()
+    new_dct = MySchema.from_json(json_str).to_dict()
+
+    assert new_dct == dct