|
1 | 1 | import graphviz |
2 | | -from aapi import Folder, SubFolder, SimpleFolder |
| 2 | + |
| 3 | +from aapi import Folder, SimpleFolder, SubFolder |
3 | 4 | from ctm_python_client.core.workflow import BaseWorkflow |
4 | 5 |
|
5 | 6 |
|
6 | 7 | def get_subgraph(obj, name, current_path): |
7 | | - if not( |
8 | | - isinstance(obj, Folder) or |
9 | | - isinstance(obj, SubFolder) or |
10 | | - isinstance(obj, SimpleFolder) |
| 8 | + if not ( |
| 9 | + isinstance(obj, Folder) |
| 10 | + or isinstance(obj, SubFolder) |
| 11 | + or isinstance(obj, SimpleFolder) |
11 | 12 | ): |
12 | 13 | return |
13 | | - sgraph = graphviz.Digraph(name, graph_attr={'label':obj.object_name}, node_attr={'shape':'oval'}) |
| 14 | + sgraph = graphviz.Digraph( |
| 15 | + name, graph_attr={"label": obj.object_name}, node_attr={"shape": "oval"} |
| 16 | + ) |
14 | 17 | for o in obj.job_list: |
15 | | - sgraph.node(f'{current_path}/{o.object_name}', o.object_name) |
| 18 | + sgraph.node(f"{current_path}/{o.object_name}", o.object_name) |
16 | 19 |
|
17 | | - for i,o in enumerate(obj.sub_folder_list): |
18 | | - sgraph.subgraph(get_subgraph(o, f'cluster_{current_path}_{i}', current_path+'/'+o.object_name)) |
| 20 | + for i, o in enumerate(obj.sub_folder_list): |
| 21 | + sgraph.subgraph( |
| 22 | + get_subgraph( |
| 23 | + o, f"cluster_{current_path}_{i}", current_path + "/" + o.object_name |
| 24 | + ) |
| 25 | + ) |
19 | 26 |
|
20 | 27 | return sgraph |
21 | 28 |
|
22 | | -def get_graph(workflow : BaseWorkflow): |
23 | | - dgraph = graphviz.Digraph(name='root', graph_attr={ |
24 | | - 'rankdir': 'LR', |
25 | | - 'compound': 'true' |
26 | | - }, |
27 | | - node_attr={'shape':'oval'} |
| 29 | + |
| 30 | +def get_graph(workflow: BaseWorkflow): |
| 31 | + dgraph = graphviz.Digraph( |
| 32 | + name="root", |
| 33 | + graph_attr={"rankdir": "LR", "compound": "true"}, |
| 34 | + node_attr={"shape": "oval"}, |
28 | 35 | ) |
29 | 36 | i = 0 |
30 | 37 | for k, v in workflow._definitions.items(): |
31 | | - if not( |
32 | | - isinstance(v, Folder) or |
33 | | - isinstance(v, SubFolder) or |
34 | | - isinstance(v, SimpleFolder) |
| 38 | + if not ( |
| 39 | + isinstance(v, Folder) |
| 40 | + or isinstance(v, SubFolder) |
| 41 | + or isinstance(v, SimpleFolder) |
35 | 42 | ): |
36 | 43 | continue |
37 | | - cluster = graphviz.Digraph(f'cluster_{i}', graph_attr={ |
38 | | - 'label': v.object_name |
39 | | - }) |
| 44 | + cluster = graphviz.Digraph(f"cluster_{i}", graph_attr={"label": v.object_name}) |
40 | 45 |
|
41 | 46 | for o in v.job_list: |
42 | | - cluster.node(f'{v.object_name}/{o.object_name}', o.object_name) |
| 47 | + cluster.node(f"{v.object_name}/{o.object_name}", o.object_name) |
43 | 48 |
|
44 | | - for i,o in enumerate(v.sub_folder_list): |
45 | | - cluster.subgraph(get_subgraph(o,f'cluster_{v.object_name}/{o.object_name}_{i}',v.object_name+'/'+o.object_name)) |
46 | | - |
| 49 | + if hasattr(v, "sub_folder_list"): |
| 50 | + for i, o in enumerate(v.sub_folder_list): |
| 51 | + cluster.subgraph( |
| 52 | + get_subgraph( |
| 53 | + o, |
| 54 | + f"cluster_{v.object_name}/{o.object_name}_{i}", |
| 55 | + v.object_name + "/" + o.object_name, |
| 56 | + ) |
| 57 | + ) |
47 | 58 |
|
48 | 59 | dgraph.subgraph(cluster) |
49 | 60 | i += 1 |
50 | 61 |
|
51 | 62 | for conn in workflow._connections.values(): |
52 | 63 | for o in conn: |
53 | | - dgraph.edge(o[0],o[1]) |
| 64 | + dgraph.edge(o[0], o[1]) |
54 | 65 |
|
55 | 66 | return dgraph |
0 commit comments