"""
Magic functions for rendering vega/vega-lite specifications
"""
__all__ = ["vega", "vegalite"]
import json
import warnings
import IPython
from IPython.core import magic_arguments
import pandas as pd
from toolz import curried
from altair.vegalite import v3 as vegalite_v3
from altair.vegalite import v4 as vegalite_v4
from altair.vega import v5 as vega_v5
try:
import yaml
YAML_AVAILABLE = True
except ImportError:
YAML_AVAILABLE = False
RENDERERS = {
"vega": {"5": vega_v5.Vega},
"vega-lite": {"3": vegalite_v3.VegaLite, "4": vegalite_v4.VegaLite},
}
TRANSFORMERS = {
"vega": {
# Vega doesn't yet have specific data transformers; use vegalite
"5": vegalite_v4.data_transformers,
},
"vega-lite": {
"3": vegalite_v3.data_transformers,
"4": vegalite_v4.data_transformers,
},
}
def _prepare_data(data, data_transformers):
"""Convert input data to data for use within schema"""
if data is None or isinstance(data, dict):
return data
elif isinstance(data, pd.DataFrame):
return curried.pipe(data, data_transformers.get())
elif isinstance(data, str):
return {"url": data}
else:
warnings.warn("data of type {} not recognized".format(type(data)))
return data
def _get_variable(name):
"""Get a variable from the notebook namespace."""
ip = IPython.get_ipython()
if ip is None:
raise ValueError(
"Magic command must be run within an IPython "
"environemnt, in which get_ipython() is defined."
)
if name not in ip.user_ns:
raise NameError(
"argument '{}' does not match the "
"name of any defined variable".format(name)
)
return ip.user_ns[name]
@magic_arguments.magic_arguments()
@magic_arguments.argument(
"data",
nargs="*",
help="local variable name of a pandas DataFrame to be used as the dataset",
)
@magic_arguments.argument("-v", "--version", dest="version", default="5")
@magic_arguments.argument("-j", "--json", dest="json", action="store_true")
def vega(line, cell):
"""Cell magic for displaying Vega visualizations in CoLab.
%%vega [name1:variable1 name2:variable2 ...] [--json] [--version='5']
Visualize the contents of the cell using Vega, optionally specifying
one or more pandas DataFrame objects to be used as the datasets.
If --json is passed, then input is parsed as json rather than yaml.
"""
args = magic_arguments.parse_argstring(vega, line)
version = args.version
assert version in RENDERERS["vega"]
Vega = RENDERERS["vega"][version]
data_transformers = TRANSFORMERS["vega"][version]
def namevar(s):
s = s.split(":")
if len(s) == 1:
return s[0], s[0]
elif len(s) == 2:
return s[0], s[1]
else:
raise ValueError("invalid identifier: '{}'".format(s))
try:
data = list(map(namevar, args.data))
except ValueError:
raise ValueError("Could not parse arguments: '{}'".format(line))
if args.json:
spec = json.loads(cell)
elif not YAML_AVAILABLE:
try:
spec = json.loads(cell)
except json.JSONDecodeError:
raise ValueError(
"%%vega: spec is not valid JSON. "
"Install pyyaml to parse spec as yaml"
)
else:
spec = yaml.load(cell, Loader=yaml.FullLoader)
if data:
spec["data"] = []
for name, val in data:
val = _get_variable(val)
prepped = _prepare_data(val, data_transformers)
prepped["name"] = name
spec["data"].append(prepped)
return Vega(spec)
@magic_arguments.magic_arguments()
@magic_arguments.argument(
"data",
nargs="?",
help="local variablename of a pandas DataFrame to be used as the dataset",
)
@magic_arguments.argument("-v", "--version", dest="version", default="4")
@magic_arguments.argument("-j", "--json", dest="json", action="store_true")
def vegalite(line, cell):
"""Cell magic for displaying vega-lite visualizations in CoLab.
%%vegalite [dataframe] [--json] [--version=3]
Visualize the contents of the cell using Vega-Lite, optionally
specifying a pandas DataFrame object to be used as the dataset.
if --json is passed, then input is parsed as json rather than yaml.
"""
args = magic_arguments.parse_argstring(vegalite, line)
version = args.version
assert version in RENDERERS["vega-lite"]
VegaLite = RENDERERS["vega-lite"][version]
data_transformers = TRANSFORMERS["vega-lite"][version]
if args.json:
spec = json.loads(cell)
elif not YAML_AVAILABLE:
try:
spec = json.loads(cell)
except json.JSONDecodeError:
raise ValueError(
"%%vegalite: spec is not valid JSON. "
"Install pyyaml to parse spec as yaml"
)
else:
spec = yaml.load(cell, Loader=yaml.FullLoader)
if args.data is not None:
data = _get_variable(args.data)
spec["data"] = _prepare_data(data, data_transformers)
return VegaLite(spec)
