Skip to content

Commit c1ea714

Browse files
williambdeanWill DeanConengmo
authored
new vega lite version (#1525)
* new vega lite version * adding support for vega-lite v5. modify the tests for checking * following the stickler-ci comments * returning int. Using None type instead * since v1 has its own format Co-authored-by: Will Dean <[email protected]> Co-authored-by: Frank <[email protected]>
1 parent 7ba08bf commit c1ea714

File tree

4 files changed

+198
-108
lines changed

4 files changed

+198
-108
lines changed

‎folium/features.py

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -231,8 +231,6 @@ def __init__(self, data, width=None, height=None,
231231

232232
def render(self, **kwargs):
233233
"""Renders the HTML representation of the element."""
234-
vegalite_major_version = self._get_vegalite_major_versions(self.data)
235-
236234
self._parent.html.add_child(Element(Template("""
237235
<div id="{{this.get_name()}}"></div>
238236
""").render(this=self, kwargs=kwargs)), name=self.get_name())
@@ -251,25 +249,40 @@ def render(self, **kwargs):
251249
</style>
252250
""").render(this=self, **kwargs)), name=self.get_name())
253251

254-
if vegalite_major_version == '1':
255-
self._embed_vegalite_v1(figure)
256-
elif vegalite_major_version == '2':
257-
self._embed_vegalite_v2(figure)
258-
elif vegalite_major_version == '3':
259-
self._embed_vegalite_v3(figure)
260-
else:
261-
# Version 2 is assumed as the default, if no version is given in the schema.
262-
self._embed_vegalite_v2(figure)
263-
264-
def _get_vegalite_major_versions(self, spec):
265-
try:
266-
schema = spec['$schema']
267-
except KeyError:
268-
major_version = None
269-
else:
270-
major_version = schema.split('/')[-1].split('.')[0].lstrip('v')
252+
embed_mapping = {
253+
1: self._embed_vegalite_v1,
254+
2: self._embed_vegalite_v2,
255+
3: self._embed_vegalite_v3,
256+
4: self._embed_vegalite_v4,
257+
5: self._embed_vegalite_v5,
258+
}
259+
260+
# Version 2 is assumed as the default, if no version is given in the schema.
261+
embed_vegalite = embed_mapping.get(self.vegalite_major_version, self._embed_vegalite_v2)
262+
embed_vegalite(figure)
263+
264+
@property
265+
def vegalite_major_version(self) -> int:
266+
if '$schema' not in self.data:
267+
return None
268+
269+
schema = self.data['$schema']
270+
271+
return int(schema.split('/')[-1].split('.')[0].lstrip('v'))
272+
273+
def _embed_vegalite_v5(self, figure):
274+
self._vega_embed()
275+
276+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm//vega@5'), name='vega')
277+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm/vega-lite@5'), name='vega-lite')
278+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm/vega-embed@6'), name='vega-embed')
279+
280+
def _embed_vegalite_v4(self, figure):
281+
self._vega_embed()
271282

272-
return major_version
283+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm//vega@5'), name='vega')
284+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm/vega-lite@4'), name='vega-lite')
285+
figure.header.add_child(JavascriptLink('https://cdn.jsdelivr.net/npm/vega-embed@6'), name='vega-embed')
273286

274287
def _embed_vegalite_v3(self, figure):
275288
self._vega_embed()

‎tests/test_features.py

Lines changed: 39 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import os
88
import warnings
99

10+
import json
11+
1012
from branca.element import Element
1113

1214
import folium
@@ -114,94 +116,43 @@ def test_color_line():
114116
m._repr_html_()
115117

116118

117-
def test_get_vegalite_major_version():
118-
spec_v2 = {'$schema': 'https://vega.github.io/schema/vega-lite/v2.6.0.json',
119-
'config': {'view': {'height': 300, 'width': 400}},
120-
'data': {'name': 'data-aac17e868e23f98b5e0830d45504be45'},
121-
'datasets': {'data-aac17e868e23f98b5e0830d45504be45': [{'folium usage': 0,
122-
'happiness': 1.0},
123-
{'folium usage': 1,
124-
'happiness': 2.718281828459045},
125-
{'folium usage': 2,
126-
'happiness': 7.38905609893065},
127-
{'folium usage': 3,
128-
'happiness': 20.085536923187668},
129-
{'folium usage': 4,
130-
'happiness': 54.598150033144236},
131-
{'folium usage': 5,
132-
'happiness': 148.4131591025766},
133-
{'folium usage': 6,
134-
'happiness': 403.4287934927351},
135-
{'folium usage': 7,
136-
'happiness': 1096.6331584284585},
137-
{'folium usage': 8,
138-
'happiness': 2980.9579870417283},
139-
{'folium usage': 9,
140-
'happiness': 8103.083927575384}]},
141-
'encoding': {'x': {'field': 'folium usage', 'type': 'quantitative'},
142-
'y': {'field': 'happiness', 'type': 'quantitative'}},
143-
'mark': 'point'}
144-
145-
vegalite_v2 = folium.features.VegaLite(spec_v2)
146-
147-
assert vegalite_v2._get_vegalite_major_versions(spec_v2) == '2'
148-
149-
spec_v1 = {'$schema': 'https://vega.github.io/schema/vega-lite/v1.3.1.json',
150-
'data': {'values': [{'folium usage': 0, 'happiness': 1.0},
151-
{'folium usage': 1, 'happiness': 2.718281828459045},
152-
{'folium usage': 2, 'happiness': 7.38905609893065},
153-
{'folium usage': 3, 'happiness': 20.085536923187668},
154-
{'folium usage': 4, 'happiness': 54.598150033144236},
155-
{'folium usage': 5, 'happiness': 148.4131591025766},
156-
{'folium usage': 6, 'happiness': 403.4287934927351},
157-
{'folium usage': 7, 'happiness': 1096.6331584284585},
158-
{'folium usage': 8, 'happiness': 2980.9579870417283},
159-
{'folium usage': 9, 'happiness': 8103.083927575384}]},
160-
'encoding': {'x': {'field': 'folium usage', 'type': 'quantitative'},
161-
'y': {'field': 'happiness', 'type': 'quantitative'}},
162-
'height': 300,
163-
'mark': 'point',
164-
'width': 400}
165-
166-
vegalite_v1 = folium.features.VegaLite(spec_v1)
167-
168-
assert vegalite_v1._get_vegalite_major_versions(spec_v1) == '1'
169-
170-
spec_no_version = {
171-
'config': {
172-
'view': {'height': 300, 'width': 400}},
173-
'data': {'name': 'data-aac17e868e23f98b5e0830d45504be45'},
174-
'datasets': {
175-
'data-aac17e868e23f98b5e0830d45504be45': [
176-
{'folium usage': 0,
177-
'happiness': 1.0},
178-
{'folium usage': 1,
179-
'happiness': 2.718281828459045},
180-
{'folium usage': 2,
181-
'happiness': 7.38905609893065},
182-
{'folium usage': 3,
183-
'happiness': 20.085536923187668},
184-
{'folium usage': 4,
185-
'happiness': 54.598150033144236},
186-
{'folium usage': 5,
187-
'happiness': 148.4131591025766},
188-
{'folium usage': 6,
189-
'happiness': 403.4287934927351},
190-
{'folium usage': 7,
191-
'happiness': 1096.6331584284585},
192-
{'folium usage': 8,
193-
'happiness': 2980.9579870417283},
194-
{'folium usage': 9,
195-
'happiness': 8103.083927575384}]},
196-
'encoding': {'x': {'field': 'folium usage', 'type': 'quantitative'},
197-
'y': {'field': 'happiness', 'type': 'quantitative'}},
198-
'mark': 'point'
199-
}
200-
201-
vegalite_no_version = folium.features.VegaLite(spec_no_version)
202-
203-
assert vegalite_no_version._get_vegalite_major_versions(spec_no_version) is None
204-
119+
@pytest.fixture
120+
def vegalite_spec(version):
121+
file_version = 'v1' if version == 1 else 'vlater'
122+
file = os.path.join(rootpath, 'vegalite_data', f'vegalite_{file_version}.json')
123+
124+
if not os.path.exists(file):
125+
raise FileNotFoundError(f'The vegalite data {file} does not exist.')
126+
127+
with open(file, 'r') as f:
128+
spec = json.load(f)
129+
130+
if version is None or '$schema' in spec:
131+
return spec
132+
133+
# Sample versions that might show up
134+
schema_version = {
135+
2: 'v2.6.0',
136+
3: 'v3.6.0',
137+
4: 'v4.6.0',
138+
5: 'v5.1.0'
139+
}[version]
140+
spec['$schema'] = f'https://vega.github.io/schema/vega-lite/{schema_version}.json'
141+
142+
return spec
143+
144+
145+
@pytest.mark.parametrize(
146+
'version',
147+
[1, 2, 3, 4, 5, None]
148+
)
149+
def test_vegalite_major_version(vegalite_spec, version):
150+
vegalite = folium.features.VegaLite(vegalite_spec)
151+
152+
if version is None:
153+
assert vegalite.vegalite_major_version is None
154+
else:
155+
assert vegalite.vegalite_major_version == version
205156

206157
# GeoJsonTooltip GeometryCollection
207158
def test_geojson_tooltip():

‎tests/vegalite_data/vegalite_v1.json

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
{
2+
"$schema": "https://vega.github.io/schema/vega-lite/v1.3.1.json",
3+
"data": {
4+
"values": [
5+
{
6+
"folium usage": 0,
7+
"happiness": 1
8+
},
9+
{
10+
"folium usage": 1,
11+
"happiness": 2.718281828459045
12+
},
13+
{
14+
"folium usage": 2,
15+
"happiness": 7.38905609893065
16+
},
17+
{
18+
"folium usage": 3,
19+
"happiness": 20.085536923187668
20+
},
21+
{
22+
"folium usage": 4,
23+
"happiness": 54.598150033144236
24+
},
25+
{
26+
"folium usage": 5,
27+
"happiness": 148.4131591025766
28+
},
29+
{
30+
"folium usage": 6,
31+
"happiness": 403.4287934927351
32+
},
33+
{
34+
"folium usage": 7,
35+
"happiness": 1096.6331584284585
36+
},
37+
{
38+
"folium usage": 8,
39+
"happiness": 2980.9579870417283
40+
},
41+
{
42+
"folium usage": 9,
43+
"happiness": 8103.083927575384
44+
}
45+
]
46+
},
47+
"encoding": {
48+
"x": {
49+
"field": "folium usage",
50+
"type": "quantitative"
51+
},
52+
"y": {
53+
"field": "happiness",
54+
"type": "quantitative"
55+
}
56+
},
57+
"height": 300,
58+
"mark": "point",
59+
"width": 400
60+
}

‎tests/vegalite_data/vegalite_vlater.json

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
{
2+
"config": {
3+
"view": {
4+
"height": 300,
5+
"width": 400
6+
}
7+
},
8+
"data": {
9+
"name": "data-aac17e868e23f98b5e0830d45504be45"
10+
},
11+
"datasets": {
12+
"data-aac17e868e23f98b5e0830d45504be45": [
13+
{
14+
"folium usage": 0,
15+
"happiness": 1
16+
},
17+
{
18+
"folium usage": 1,
19+
"happiness": 2.718281828459045
20+
},
21+
{
22+
"folium usage": 2,
23+
"happiness": 7.38905609893065
24+
},
25+
{
26+
"folium usage": 3,
27+
"happiness": 20.085536923187668
28+
},
29+
{
30+
"folium usage": 4,
31+
"happiness": 54.598150033144236
32+
},
33+
{
34+
"folium usage": 5,
35+
"happiness": 148.4131591025766
36+
},
37+
{
38+
"folium usage": 6,
39+
"happiness": 403.4287934927351
40+
},
41+
{
42+
"folium usage": 7,
43+
"happiness": 1096.6331584284585
44+
},
45+
{
46+
"folium usage": 8,
47+
"happiness": 2980.9579870417283
48+
},
49+
{
50+
"folium usage": 9,
51+
"happiness": 8103.083927575384
52+
}
53+
]
54+
},
55+
"encoding": {
56+
"x": {
57+
"field": "folium usage",
58+
"type": "quantitative"
59+
},
60+
"y": {
61+
"field": "happiness",
62+
"type": "quantitative"
63+
}
64+
},
65+
"mark": "point"
66+
}

0 commit comments

Comments
 (0)