|
""" |
|
Magic functions for rendering vega-lite specifications |
|
""" |
|
__all__ = ["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 v5 as vegalite_v5 |
|
|
|
try: |
|
import yaml |
|
|
|
YAML_AVAILABLE = True |
|
except ImportError: |
|
YAML_AVAILABLE = False |
|
|
|
|
|
RENDERERS = { |
|
"vega-lite": { |
|
"5": vegalite_v5.VegaLite, |
|
}, |
|
} |
|
|
|
|
|
TRANSFORMERS = { |
|
"vega-lite": { |
|
"5": vegalite_v5.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)), stacklevel=1) |
|
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 " |
|
"environment, 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 variablename of a pandas DataFrame to be used as the dataset", |
|
) |
|
@magic_arguments.argument("-v", "--version", dest="version", default="v5") |
|
@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='v5'] |
|
|
|
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) |
|
existing_versions = {"v5": "5"} |
|
version = existing_versions[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 as err: |
|
raise ValueError( |
|
"%%vegalite: spec is not valid JSON. " |
|
"Install pyyaml to parse spec as yaml" |
|
) from err |
|
else: |
|
spec = yaml.load(cell, Loader=yaml.SafeLoader) |
|
|
|
if args.data is not None: |
|
data = _get_variable(args.data) |
|
spec["data"] = _prepare_data(data, data_transformers) |
|
|
|
return VegaLite(spec) |
|
|
