Repository: altair-viz/altair
Pull request #465 (2018-02-23 17:33:58)

MAINT: add  property to top-level objects on to_dict().
This will make them automatically render correctly when exported to the vega editor.


Patch:
diff --git a/altair/vegalite/v1/api.py b/altair/vegalite/v1/api.py
index 52637056..75955ba1 100644
--- a/altair/vegalite/v1/api.py
+++ b/altair/vegalite/v1/api.py
@@ -16,6 +16,9 @@ from ...utils import (infer_vegalite_type, parse_shorthand_plus_data,
 from .display import renderers
 
 
+SCHEMA_URL = "https://vega.github.io/schema/vega-lite/v1.json"
+
+
 def _get_channels_mapping():
     mapping = {}
     for attr in dir(channels):
@@ -80,8 +83,6 @@ def _get_channels_mapping():
 
 class TopLevelMixin(object):
     def _prepare_data(self):
-        # TODO: it's a bit weird that to_dict, via _prepare_data,
-        #       modifies the object. Should we create a copy first?
         if isinstance(self.data, (dict, core.Data)):
             pass
         elif isinstance(self.data, pd.DataFrame):
@@ -90,11 +91,25 @@ class TopLevelMixin(object):
             self.data = core.Data(url=self.data)
 
     def to_dict(self, *args, **kwargs):
+        # TODO: it's a bit weird that to_dict modifies the object.
+        #       Should we create a copy first?
         original_data = getattr(self, 'data', Undefined)
         self._prepare_data()
+        # We make use of two context markers:
+        # - 'data' points to the data that should be referenced for column type
+        #   inference.
+        # - 'toplevel' 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', {})
         if 'data' not in context and original_data is not Undefined:
             context['data'] = original_data
+        if context.get('top_level', True):
+            # since this is top-level we add $schema if it's missing
+            if '$schema' not in self._kwds:
+                self._kwds['$schema'] = SCHEMA_URL
+            # subschemas below this one are not top-level
+            context['top_level'] = False
         kwargs['context'] = context
         return super(TopLevelMixin, self).to_dict(*args, **kwargs)
 
diff --git a/altair/vegalite/v2/api.py b/altair/vegalite/v2/api.py
index 429a2ba3..6b58e017 100644
--- a/altair/vegalite/v2/api.py
+++ b/altair/vegalite/v2/api.py
@@ -10,6 +10,9 @@ from ...utils import infer_vegalite_type, parse_shorthand_plus_data, use_signatu
 from .display import renderers
 
 
+SCHEMA_URL = "https://vega.github.io/schema/vega-lite/v2.json"
+
+
 def _get_channels_mapping():
     mapping = {}
     for attr in dir(channels):
@@ -21,8 +24,6 @@ def _get_channels_mapping():
 
 class TopLevelMixin(object):
     def _prepare_data(self):
-        # TODO: it's a bit weird that to_dict, via _prepare_data,
-        #       modifies the object. Should we create a copy first?
         if isinstance(self.data, (dict, core.Data, core.InlineData,
                                   core.UrlData, core.NamedData)):
             pass
@@ -32,11 +33,25 @@ class TopLevelMixin(object):
             self.data = core.UrlData(self.data)
 
     def to_dict(self, *args, **kwargs):
+        # TODO: it's a bit weird that to_dict modifies the object.
+        #       Should we create a copy first?
         original_data = getattr(self, 'data', Undefined)
         self._prepare_data()
+
+        # We make use of two context markers:
+        # - 'data' points to the data that should be referenced for column type
+        #   inference.
+        # - 'toplevel' 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', {})
         if 'data' not in context and original_data is not Undefined:
             context['data'] = original_data
+        if context.get('top_level', True):
+            # since this is top-level we add $schema if it's missing
+            if '$schema' not in self._kwds:
+                self._kwds['$schema'] = SCHEMA_URL
+            # subschemas below this one are not top-level
+            context['top_level'] = False
         kwargs['context'] = context
         return super(TopLevelMixin, self).to_dict(*args, **kwargs)
 


Test patch:
diff --git a/altair/vegalite/tests/test_common.py b/altair/vegalite/tests/test_common.py
index b2bcbe43..4034c62d 100644
--- a/altair/vegalite/tests/test_common.py
+++ b/altair/vegalite/tests/test_common.py
@@ -28,10 +28,21 @@ def test_basic_chart_to_dict(alt, basic_spec):
         color = 'color:N'
     )
     dct = chart.to_dict()
+
+    # schema should be in the top level
+    assert dct.pop('$schema').startswith('http')
+
+    # remainder of spec should match the basic spec
     assert dct == basic_spec
 
 
 @pytest.mark.parametrize('alt', [v1, v2])
 def test_basic_chart_from_dict(alt, basic_spec):
     chart = alt.Chart.from_dict(basic_spec)
-    assert chart.to_dict() == basic_spec
+    dct = chart.to_dict()
+
+    # schema should be in the top level
+    assert dct.pop('$schema').startswith('http')
+
+    # remainder of spec should match the basic spec
+    assert dct == basic_spec