def test_sankey_definition_as_script_with_partitions():
nodes = {
'a':
ProcessGroup(selection=['a1', 'a2']),
'b':
ProcessGroup(selection=['b1']),
'c':
ProcessGroup(selection=['c1', 'c2'],
partition=Partition.Simple('process', ['c1', 'c2'])),
'via':
Waypoint(partition=Partition.Simple('material', ['m', 'n'])),
}
bundles = [
Bundle('a', 'c', waypoints=['via']),
Bundle('b', 'c', waypoints=['via']),
]
ordering = [[['a', 'b']], [['via']], [['c']]]
sdd = SankeyDefinition(nodes,
bundles,
ordering,
flow_partition=Partition.Simple(
'material', ['m', 'n']))
code = sdd.to_code()
# Check roundtrip
ctx = {}
exec(code, ctx)
assert ctx["sdd"] == sdd
def test_sankey_definition_as_script():
nodes = {
'a': ProcessGroup(selection=['a1']),
'b': ProcessGroup(selection=['b1']),
'waypoint': Waypoint(),
}
ordering = [['a'], ['waypoint'], ['b']]
bundles = [Bundle('a', 'b')]
sdd = SankeyDefinition(nodes, bundles, ordering)
code = sdd.to_code()
assert code == dedent("""
from floweaver import (
ProcessGroup,
Waypoint,
Partition,
Group,
Elsewhere,
Bundle,
SankeyDefinition,
)
nodes = {
'a': ProcessGroup(selection=['a1']),
'b': ProcessGroup(selection=['b1']),
'waypoint': Waypoint(),
}
ordering = [
[['a']],
[['waypoint']],
[['b']],
]
bundles = [
Bundle(source='a', target='b'),
]
sdd = SankeyDefinition(nodes, bundles, ordering)
""")
# Check roundtrip
ctx = {}
exec(code, ctx)
assert ctx["sdd"] == sdd