1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
import json
import pytest
try:
from IPython import InteractiveShell
IPYTHON_AVAILABLE = True
except ImportError:
IPYTHON_AVAILABLE = False
pass
from altair.vegalite.v5 import VegaLite
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
)
)
VEGALITE_SPEC = {
"$schema": "https://vega.github.io/schema/vega-lite/v5.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": {"type": "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
|
