MAINT: use standard form for magics

saulshanabrook · Apr 11, 2018 · 033e632 · 033e632
1 parent 0369c7e
commit 033e632
Show file tree

Hide file tree

Showing 4 changed files with 2,230 additions and 33 deletions.
diff --git a/altair/__init__.py b/altair/__init__.py
@@ -2,3 +2,8 @@
 __version__ = '2.0.0dev'
 
 from .vegalite import *
+
+def load_ipython_extension(ipython):
+    from ._magics import vega, vegalite
+    ipython.register_magic_function(vega, 'cell')
+    ipython.register_magic_function(vegalite, 'cell')
diff --git a/altair/magics.py → altair/_magics.py b/altair/magics.py → altair/_magics.py
@@ -40,12 +40,12 @@
 TRANSFORMERS = {
   'vega': {
       # Vega doesn't yet have specific data transformers; use vegalite
-      '2': vegalite_v1.data.data_transformers,
-      '3': vegalite_v2.data.data_transformers,
+      '2': vegalite_v1.data_transformers,
+      '3': vegalite_v2.data_transformers,
   },
   'vega-lite': {
-      '1': vegalite_v1.data.data_transformers,
-      '2': vegalite_v2.data.data_transformers,
+      '1': vegalite_v1.data_transformers,
+      '2': vegalite_v2.data_transformers,
   }
 }
 
@@ -75,7 +75,6 @@ def _get_variable(name):
     return ip.user_ns[name]
 
 
-@magic.register_cell_magic
 @magic_arguments.magic_arguments()
 @magic_arguments.argument(
     'data',
@@ -119,7 +118,7 @@ def namevar(s):
     elif not YAML_AVAILABLE:
         try:
             spec = json.loads(cell)
-        except json.JSONDecodeError:
+        except JSONDecodeError:
             raise ValueError("%%vega: spec is not valid JSON. "
                              "Install pyyaml to parse spec as yaml")
     else:
@@ -136,7 +135,6 @@ def namevar(s):
     return Vega(spec)
 
 
-@magic.register_cell_magic
 @magic_arguments.magic_arguments()
 @magic_arguments.argument(
     'data',
@@ -165,7 +163,7 @@ def vegalite(line, cell):
     elif not YAML_AVAILABLE:
         try:
             spec = json.loads(cell)
-        except json.JSONDecodeError:
+        except JSONDecodeError:
             raise ValueError("%%vegalite: spec is not valid JSON. "
                              "Install pyyaml to parse spec as yaml")
     else:

diff --git a/altair/tests/test_magics.py b/altair/tests/test_magics.py
@@ -6,7 +6,7 @@
 from altair.vega.v3 import Vega
 
 _ipshell = InteractiveShell.instance()
-_ipshell.run_cell('from altair.magics import vega, vegalite')
+_ipshell.run_cell('%load_ext altair')
 
 DATA_RECORDS = [{'amount': 28, 'category': 'A'},
                 {'amount': 55, 'category': 'B'},
@@ -78,7 +78,7 @@ def test_vegalite_magic_data_included():
     result = _ipshell.run_cell('%%vegalite\n' + json.dumps(VEGALITE_SPEC))
     assert isinstance(result.result, VegaLite)
     assert VEGALITE_SPEC == result.result.spec
-    
+
 
 def test_vegalite_magic_json_flag():
     result = _ipshell.run_cell('%%vegalite --json\n'