test_sankey.py

#!/usr/bin/env python
# coding=utf-8
from __future__ import unicode_literals
import sys
import os

from pyecharts import Sankey

PY2 = sys.version_info[0] == 2


def test_sankey_default():
    nodes = [
        {'name': 'category1'}, {'name': 'category2'}, {'name': 'category3'},
        {'name': 'category4'}, {'name': 'category5'}, {'name': 'category6'},
    ]

    links = [
        {'source': 'category1', 'target': 'category2', 'value': 10},
        {'source': 'category2', 'target': 'category3', 'value': 15},
        {'source': 'category3', 'target': 'category4', 'value': 20},
        {'source': 'category5', 'target': 'category6', 'value': 25}
    ]
    sankey = Sankey("桑基图示例", width=1200, height=600)
    sankey.add("sankey", nodes, links, line_opacity=0.2,
               line_curve=0.5, line_color='source', is_label_show=True,
               label_pos='right')
    sankey.render()


def test_sankey_official_data():
    import json
    if PY2:
        import codecs
        with codecs.open(os.path.join("..", "json", "energy.json"), "rb") as f:
            j = json.load(f)
    else:
        with open(os.path.join("..", "json", "energy.json"),
                  "r", encoding="utf-8") as f:
            j = json.load(f)
    sankey = Sankey("桑基图示例", width=1200, height=600)
    sankey.add("sankey", nodes=j['nodes'], links=j['links'], line_opacity=0.2,
               line_curve=0.5, line_color='source',
               is_label_show=True, label_pos='right')
    sankey.render()