venv/lib/python3.11/site-packages/altair/_magics.py

"""
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)