runtime/Lib/site-packages/altair/tests/test_magics.py

import json
import pytest

try:
    from IPython import InteractiveShell

    IPYTHON_AVAILABLE = True
except ImportError:
    IPYTHON_AVAILABLE = False
    pass

from altair.vegalite.v4 import VegaLite
from altair.vega.v5 import Vega


DATA_RECORDS = [
    {"amount": 28, "category": "A"},
    {"amount": 55, "category": "B"},
    {"amount": 43, "category": "C"},
    {"amount": 91, "category": "D"},
    {"amount": 81, "category": "E"},
    {"amount": 53, "category": "F"},
    {"amount": 19, "category": "G"},
    {"amount": 87, "category": "H"},
]

if IPYTHON_AVAILABLE:
    _ipshell = InteractiveShell.instance()
    _ipshell.run_cell("%load_ext altair")
    _ipshell.run_cell(
        """
import pandas as pd
table = pd.DataFrame.from_records({})
the_data = table
""".format(
            DATA_RECORDS
        )
    )


VEGA_SPEC = {
    "$schema": "https://vega.github.io/schema/vega/v5.json",
    "axes": [
        {"orient": "bottom", "scale": "xscale"},
        {"orient": "left", "scale": "yscale"},
    ],
    "data": [{"name": "table", "values": DATA_RECORDS}],
    "height": 200,
    "marks": [
        {
            "encode": {
                "enter": {
                    "width": {"band": 1, "scale": "xscale"},
                    "x": {"field": "category", "scale": "xscale"},
                    "y": {"field": "amount", "scale": "yscale"},
                    "y2": {"scale": "yscale", "value": 0},
                },
                "hover": {"fill": {"value": "red"}},
                "update": {"fill": {"value": "steelblue"}},
            },
            "from": {"data": "table"},
            "type": "rect",
        },
        {
            "encode": {
                "enter": {
                    "align": {"value": "center"},
                    "baseline": {"value": "bottom"},
                    "fill": {"value": "#333"},
                },
                "update": {
                    "fillOpacity": [
                        {"test": "datum === tooltip", "value": 0},
                        {"value": 1},
                    ],
                    "text": {"signal": "tooltip.amount"},
                    "x": {"band": 0.5, "scale": "xscale", "signal": "tooltip.category"},
                    "y": {"offset": -2, "scale": "yscale", "signal": "tooltip.amount"},
                },
            },
            "type": "text",
        },
    ],
    "padding": 5,
    "scales": [
        {
            "domain": {"data": "table", "field": "category"},
            "name": "xscale",
            "padding": 0.05,
            "range": "width",
            "round": True,
            "type": "band",
        },
        {
            "domain": {"data": "table", "field": "amount"},
            "name": "yscale",
            "nice": True,
            "range": "height",
        },
    ],
    "signals": [
        {
            "name": "tooltip",
            "on": [
                {"events": "rect:mouseover", "update": "datum"},
                {"events": "rect:mouseout", "update": "{}"},
            ],
            "value": {},
        }
    ],
    "width": 400,
}


VEGALITE_SPEC = {
    "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
    "data": {"values": DATA_RECORDS},
    "description": "A simple bar chart with embedded data.",
    "encoding": {
        "x": {"field": "category", "type": "ordinal"},
        "y": {"field": "amount", "type": "quantitative"},
    },
    "mark": "bar",
}


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
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


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vegalite_magic_json_flag():
    result = _ipshell.run_cell("%%vegalite --json\n" + json.dumps(VEGALITE_SPEC))
    assert isinstance(result.result, VegaLite)
    assert VEGALITE_SPEC == result.result.spec


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vegalite_magic_pandas_data():
    spec = {key: val for key, val in VEGALITE_SPEC.items() if key != "data"}
    result = _ipshell.run_cell("%%vegalite table\n" + json.dumps(spec))
    assert isinstance(result.result, VegaLite)
    assert VEGALITE_SPEC == result.result.spec


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vega_magic_data_included():
    result = _ipshell.run_cell("%%vega\n" + json.dumps(VEGA_SPEC))
    assert isinstance(result.result, Vega)
    assert VEGA_SPEC == result.result.spec


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vega_magic_json_flag():
    result = _ipshell.run_cell("%%vega --json\n" + json.dumps(VEGA_SPEC))
    assert isinstance(result.result, Vega)
    assert VEGA_SPEC == result.result.spec


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vega_magic_pandas_data():
    spec = {key: val for key, val in VEGA_SPEC.items() if key != "data"}
    result = _ipshell.run_cell("%%vega table\n" + json.dumps(spec))
    assert isinstance(result.result, Vega)
    assert VEGA_SPEC == result.result.spec


@pytest.mark.skipif(not IPYTHON_AVAILABLE, reason="requires ipython")
def test_vega_magic_pandas_data_renamed():
    spec = {key: val for key, val in VEGA_SPEC.items() if key != "data"}
    result = _ipshell.run_cell("%%vega table:the_data\n" + json.dumps(spec))
    assert isinstance(result.result, Vega)
    assert VEGA_SPEC == result.result.spec