From fe5ffefbcc16b47f5af65f57517c4af6628698dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20Jan=C3=9Fen?= Date: Sat, 12 Apr 2025 11:29:08 +0200 Subject: [PATCH] Simplify updating the format by defining the labels only once --- .../src/python_workflow_definition/aiida.py | 60 +++++++++-------- .../python_workflow_definition/executorlib.py | 18 ++++-- .../src/python_workflow_definition/jobflow.py | 64 +++++++++++++------ .../python_workflow_definition/purepython.py | 30 ++++++--- .../python_workflow_definition/pyiron_base.py | 56 +++++++++------- .../src/python_workflow_definition/shared.py | 22 +++++-- 6 files changed, 162 insertions(+), 88 deletions(-) diff --git a/python_workflow_definition/src/python_workflow_definition/aiida.py b/python_workflow_definition/src/python_workflow_definition/aiida.py index bfb51c0..5a4d431 100644 --- a/python_workflow_definition/src/python_workflow_definition/aiida.py +++ b/python_workflow_definition/src/python_workflow_definition/aiida.py @@ -7,7 +7,15 @@ from aiida_workgraph import WorkGraph, task from aiida_workgraph.socket import TaskSocketNamespace -from python_workflow_definition.shared import convert_nodes_list_to_dict +from python_workflow_definition.shared import ( + convert_nodes_list_to_dict, + NODES_LABEL, + EDGES_LABEL, + SOURCE_LABEL, + SOURCE_PORT_LABEL, + TARGET_LABEL, + TARGET_PORT_LABEL, +) def load_workflow_json(file_name): @@ -17,7 +25,7 @@ def load_workflow_json(file_name): wg = WorkGraph() task_name_mapping = {} - for id, identifier in convert_nodes_list_to_dict(nodes_list=data["nodes"]).items(): + for id, identifier in convert_nodes_list_to_dict(nodes_list=data[NODES_LABEL]).items(): if isinstance(identifier, str) and "." in identifier: p, m = identifier.rsplit(".", 1) mod = import_module(p) @@ -32,33 +40,33 @@ def load_workflow_json(file_name): task_name_mapping[id] = data_node # add links - for link in data["edges"]: - to_task = task_name_mapping[str(link["target"])] + for link in data[EDGES_LABEL]: + to_task = task_name_mapping[str(link[TARGET_LABEL])] # if the input is not exit, it means we pass the data into to the kwargs # in this case, we add the input socket - if link["targetPort"] not in to_task.inputs: - to_socket = to_task.add_input( "workgraph.any", name=link["targetPort"]) + if link[TARGET_PORT_LABEL] not in to_task.inputs: + to_socket = to_task.add_input( "workgraph.any", name=link[TARGET_PORT_LABEL]) else: - to_socket = to_task.inputs[link["targetPort"]] - from_task = task_name_mapping[str(link["source"])] + to_socket = to_task.inputs[link[TARGET_PORT_LABEL]] + from_task = task_name_mapping[str(link[SOURCE_LABEL])] if isinstance(from_task, orm.Data): to_socket.value = from_task else: try: - if link["sourcePort"] is None: - link["sourcePort"] = "result" + if link[SOURCE_PORT_LABEL] is None: + link[SOURCE_PORT_LABEL] = "result" # because we are not define the outputs explicitly during the pythonjob creation # we add it here, and assume the output exit - if link["sourcePort"] not in from_task.outputs: + if link[SOURCE_PORT_LABEL] not in from_task.outputs: # if str(link["sourcePort"]) not in from_task.outputs: from_socket = from_task.add_output( "workgraph.any", - name=link["sourcePort"], + name=link[SOURCE_PORT_LABEL], # name=str(link["sourcePort"]), metadata={"is_function_output": True}, ) else: - from_socket = from_task.outputs[link["sourcePort"]] + from_socket = from_task.outputs[link[SOURCE_PORT_LABEL]] wg.add_link(from_socket, to_socket) except Exception as e: @@ -68,7 +76,7 @@ def load_workflow_json(file_name): def write_workflow_json(wg, file_name): - data = {"nodes": [], "edges": []} + data = {NODES_LABEL: [], EDGES_LABEL: []} node_name_mapping = {} data_node_name_mapping = {} i = 0 @@ -78,7 +86,7 @@ def write_workflow_json(wg, file_name): callable_name = executor["callable_name"] callable_name = f"{executor['module_path']}.{callable_name}" - data["nodes"].append({"id": i, "function": callable_name}) + data[NODES_LABEL].append({"id": i, "function": callable_name}) i += 1 for link in wg.links: @@ -86,11 +94,11 @@ def write_workflow_json(wg, file_name): # if the from socket is the default result, we set it to None if link_data["from_socket"] == "result": link_data["from_socket"] = None - link_data["target"] = node_name_mapping[link_data.pop("to_node")] - link_data["targetPort"] = link_data.pop("to_socket") - link_data["source"] = node_name_mapping[link_data.pop("from_node")] - link_data["sourcePort"] = link_data.pop("from_socket") - data["edges"].append(link_data) + link_data[TARGET_LABEL] = node_name_mapping[link_data.pop("to_node")] + link_data[TARGET_PORT_LABEL] = link_data.pop("to_socket") + link_data[SOURCE_LABEL] = node_name_mapping[link_data.pop("from_node")] + link_data[SOURCE_PORT_LABEL] = link_data.pop("from_socket") + data[EDGES_LABEL].append(link_data) for node in wg.tasks: for input in node.inputs: @@ -107,17 +115,17 @@ def write_workflow_json(wg, file_name): raw_value.pop("node_type", None) else: raw_value = input.value.value - data["nodes"].append({"id": i, "value": raw_value}) + data[NODES_LABEL].append({"id": i, "value": raw_value}) input_node_name = i data_node_name_mapping[input.value.uuid] = input_node_name i += 1 else: input_node_name = data_node_name_mapping[input.value.uuid] - data["edges"].append({ - "target": node_name_mapping[node.name], - "targetPort": input._name, - "source": input_node_name, - "sourcePort": None + data[EDGES_LABEL].append({ + TARGET_LABEL: node_name_mapping[node.name], + TARGET_PORT_LABEL: input._name, + SOURCE_LABEL: input_node_name, + SOURCE_PORT_LABEL: None }) with open(file_name, "w") as f: # json.dump({"nodes": data[], "edges": edges_new_lst}, f) diff --git a/python_workflow_definition/src/python_workflow_definition/executorlib.py b/python_workflow_definition/src/python_workflow_definition/executorlib.py index 819fd71..69ff17d 100644 --- a/python_workflow_definition/src/python_workflow_definition/executorlib.py +++ b/python_workflow_definition/src/python_workflow_definition/executorlib.py @@ -3,7 +3,17 @@ from inspect import isfunction -from python_workflow_definition.shared import get_dict, get_list, get_kwargs, get_source_handles, convert_nodes_list_to_dict +from python_workflow_definition.shared import ( + get_dict, + get_list, + get_kwargs, + get_source_handles, + convert_nodes_list_to_dict, + NODES_LABEL, + EDGES_LABEL, + SOURCE_LABEL, + SOURCE_PORT_LABEL, +) from python_workflow_definition.purepython import resort_total_lst, group_edges @@ -12,7 +22,7 @@ def get_item(obj, key): def _get_value(result_dict, nodes_new_dict, link_dict, exe): - source, source_handle = link_dict["source"], link_dict["sourcePort"] + source, source_handle = link_dict[SOURCE_LABEL], link_dict[SOURCE_PORT_LABEL] if source in result_dict.keys(): result = result_dict[source] elif source in nodes_new_dict.keys(): @@ -29,10 +39,10 @@ def load_workflow_json(file_name, exe): with open(file_name, "r") as f: content = json.load(f) - edges_new_lst = content["edges"] + edges_new_lst = content[EDGES_LABEL] nodes_new_dict = {} - for k, v in convert_nodes_list_to_dict(nodes_list=content["nodes"]).items(): + for k, v in convert_nodes_list_to_dict(nodes_list=content[NODES_LABEL]).items(): if isinstance(v, str) and "." in v: p, m = v.rsplit('.', 1) mod = import_module(p) diff --git a/python_workflow_definition/src/python_workflow_definition/jobflow.py b/python_workflow_definition/src/python_workflow_definition/jobflow.py index a5ebdce..98244e9 100644 --- a/python_workflow_definition/src/python_workflow_definition/jobflow.py +++ b/python_workflow_definition/src/python_workflow_definition/jobflow.py @@ -5,7 +5,19 @@ import numpy as np from jobflow import job, Flow -from python_workflow_definition.shared import get_dict, get_list, get_kwargs, get_source_handles, convert_nodes_list_to_dict +from python_workflow_definition.shared import ( + get_dict, + get_list, + get_kwargs, + get_source_handles, + convert_nodes_list_to_dict, + NODES_LABEL, + EDGES_LABEL, + SOURCE_LABEL, + SOURCE_PORT_LABEL, + TARGET_LABEL, + TARGET_PORT_LABEL, +) def _get_function_dict(flow): @@ -26,9 +38,19 @@ def _get_nodes_dict(function_dict): def _get_edge_from_dict(target, key, value_dict, nodes_mapping_dict): if len(value_dict['attributes']) == 1: - return {"target": target, "targetPort": key, "source": nodes_mapping_dict[value_dict["uuid"]], "sourcePort": value_dict["attributes"][0][1]} + return { + TARGET_LABEL: target, + TARGET_PORT_LABEL: key, + SOURCE_LABEL: nodes_mapping_dict[value_dict["uuid"]], + SOURCE_PORT_LABEL: value_dict["attributes"][0][1], + } else: - return {"target": target, "targetPort": key, "source": nodes_mapping_dict[value_dict["uuid"]], "sourcePort": None} + return { + TARGET_LABEL: target, + TARGET_PORT_LABEL: key, + SOURCE_LABEL: nodes_mapping_dict[value_dict["uuid"]], + SOURCE_PORT_LABEL: None, + } def _get_edges_and_extend_nodes(flow_dict, nodes_mapping_dict, nodes_dict): @@ -59,8 +81,8 @@ def _get_edges_and_extend_nodes(flow_dict, nodes_mapping_dict, nodes_dict): nodes_dict[node_index] = vt else: node_index = {str(tv): tk for tk, tv in nodes_dict.items()}[str(vt)] - edges_lst.append({"target": node_dict_index, "targetPort": kt, "source": node_index, "sourcePort": None}) - edges_lst.append({"target": nodes_mapping_dict[job["uuid"]], "targetPort": k, "source": node_dict_index, "sourcePort": None}) + edges_lst.append({TARGET_LABEL: node_dict_index, TARGET_PORT_LABEL: kt, SOURCE_LABEL: node_index, SOURCE_PORT_LABEL: None}) + edges_lst.append({TARGET_LABEL: nodes_mapping_dict[job["uuid"]], TARGET_PORT_LABEL: k, SOURCE_LABEL: node_dict_index, SOURCE_PORT_LABEL: None}) elif isinstance(v, list) and any([isinstance(el, dict) and "@module" in el and "@class" in el and "@version" in el for el in v]): node_list_index = len(nodes_dict) nodes_dict[node_list_index] = get_list @@ -78,15 +100,15 @@ def _get_edges_and_extend_nodes(flow_dict, nodes_mapping_dict, nodes_dict): nodes_dict[node_index] = vt else: node_index = {str(tv): tk for tk, tv in nodes_dict.items()}[str(vt)] - edges_lst.append({"target": node_list_index, "targetPort": kt, "source": node_index, "sourcePort": None}) - edges_lst.append({"target": nodes_mapping_dict[job["uuid"]], "targetPort": k, "source": node_list_index, "sourcePort": None}) + edges_lst.append({TARGET_LABEL: node_list_index, TARGET_PORT_LABEL: kt, SOURCE_LABEL: node_index, SOURCE_PORT_LABEL: None}) + edges_lst.append({TARGET_LABEL: nodes_mapping_dict[job["uuid"]], TARGET_PORT_LABEL: k, SOURCE_LABEL: node_list_index, SOURCE_PORT_LABEL: None}) else: if v not in nodes_dict.values(): node_index = len(nodes_dict) nodes_dict[node_index] = v else: node_index = {tv: tk for tk, tv in nodes_dict.items()}[v] - edges_lst.append({"target": nodes_mapping_dict[job["uuid"]], "targetPort": k, "source": node_index, "sourcePort": None}) + edges_lst.append({TARGET_LABEL: nodes_mapping_dict[job["uuid"]], TARGET_PORT_LABEL: k, SOURCE_LABEL: node_index, SOURCE_PORT_LABEL: None}) return edges_lst, nodes_dict @@ -99,7 +121,7 @@ def _resort_total_lst(total_dict, nodes_dict): for ind in sorted(total_dict.keys()): connect = total_dict[ind] if ind not in ordered_lst: - source_lst = [sd["source"] for sd in connect.values()] + source_lst = [sd[SOURCE_LABEL] for sd in connect.values()] if all([s in ordered_lst or s in nodes_without_dep_lst for s in source_lst]): ordered_lst.append(ind) total_new_dict[ind] = connect @@ -109,11 +131,11 @@ def _resort_total_lst(total_dict, nodes_dict): def _group_edges(edges_lst): total_dict = {} for ed_major in edges_lst: - target_id = ed_major["target"] + target_id = ed_major[TARGET_LABEL] tmp_lst = [] if target_id not in total_dict.keys(): for ed in edges_lst: - if target_id == ed["target"]: + if target_id == ed[TARGET_LABEL]: tmp_lst.append(ed) total_dict[target_id] = get_kwargs(lst=tmp_lst) return total_dict @@ -139,8 +161,8 @@ def get_attr_helper(obj, source_handle): else: fn = job(method=v) kwargs = { - kw: input_dict[vw["source"]] if vw["source"] in input_dict else get_attr_helper( - obj=memory_dict[vw["source"]], source_handle=vw["sourcePort"]) + kw: input_dict[vw[SOURCE_LABEL]] if vw[SOURCE_LABEL] in input_dict else get_attr_helper( + obj=memory_dict[vw[SOURCE_LABEL]], source_handle=vw[SOURCE_PORT_LABEL]) for kw, vw in total_dict[k].items() } memory_dict[k] = fn(**kwargs) @@ -159,21 +181,21 @@ def load_workflow_json(file_name): content = json.load(f) edges_new_lst = [] - for edge in content["edges"]: - if edge["sourcePort"] is None: + for edge in content[EDGES_LABEL]: + if edge[SOURCE_PORT_LABEL] is None: edges_new_lst.append(edge) else: edges_new_lst.append( { - "target": edge["target"], - "targetPort": edge["targetPort"], - "source": edge["source"], - "sourcePort": str(edge["sourcePort"]), + TARGET_LABEL: edge[TARGET_LABEL], + TARGET_PORT_LABEL: edge[TARGET_PORT_LABEL], + SOURCE_LABEL: edge[SOURCE_LABEL], + SOURCE_PORT_LABEL: str(edge[SOURCE_PORT_LABEL]), } ) nodes_new_dict = {} - for k, v in convert_nodes_list_to_dict(nodes_list=content["nodes"]).items(): + for k, v in convert_nodes_list_to_dict(nodes_list=content[NODES_LABEL]).items(): if isinstance(v, str) and "." in v: p, m = v.rsplit('.', 1) mod = import_module(p) @@ -214,4 +236,4 @@ def write_workflow_json(flow, file_name="workflow.json"): nodes_store_lst.append({"id": k, "value": v}) with open(file_name, "w") as f: - json.dump({"nodes": nodes_store_lst, "edges": edges_lst}, f) + json.dump({NODES_LABEL: nodes_store_lst, EDGES_LABEL: edges_lst}, f) diff --git a/python_workflow_definition/src/python_workflow_definition/purepython.py b/python_workflow_definition/src/python_workflow_definition/purepython.py index 86ac21c..cb32496 100644 --- a/python_workflow_definition/src/python_workflow_definition/purepython.py +++ b/python_workflow_definition/src/python_workflow_definition/purepython.py @@ -3,7 +3,19 @@ from inspect import isfunction -from python_workflow_definition.shared import get_dict, get_list, get_kwargs, get_source_handles, convert_nodes_list_to_dict +from python_workflow_definition.shared import ( + get_dict, + get_list, + get_kwargs, + get_source_handles, + convert_nodes_list_to_dict, + NODES_LABEL, + EDGES_LABEL, + SOURCE_LABEL, + SOURCE_PORT_LABEL, + TARGET_LABEL, + TARGET_PORT_LABEL, +) def resort_total_lst(total_lst, nodes_dict): @@ -13,7 +25,7 @@ def resort_total_lst(total_lst, nodes_dict): while len(total_new_lst) < len(total_lst): for ind, connect in total_lst: if ind not in ordered_lst: - source_lst = [sd["source"] for sd in connect.values()] + source_lst = [sd[SOURCE_LABEL] for sd in connect.values()] if all([s in ordered_lst or s in nodes_without_dep_lst for s in source_lst]): ordered_lst.append(ind) total_new_lst.append([ind, connect]) @@ -21,22 +33,22 @@ def resort_total_lst(total_lst, nodes_dict): def group_edges(edges_lst): - edges_sorted_lst = sorted(edges_lst, key=lambda x: x["target"], reverse=True) + edges_sorted_lst = sorted(edges_lst, key=lambda x: x[TARGET_LABEL], reverse=True) total_lst, tmp_lst = [], [] - target_id = edges_sorted_lst[0]["target"] + target_id = edges_sorted_lst[0][TARGET_LABEL] for ed in edges_sorted_lst: - if target_id == ed["target"]: + if target_id == ed[TARGET_LABEL]: tmp_lst.append(ed) else: total_lst.append((target_id, get_kwargs(lst=tmp_lst))) - target_id = ed["target"] + target_id = ed[TARGET_LABEL] tmp_lst = [ed] total_lst.append((target_id, get_kwargs(lst=tmp_lst))) return total_lst def _get_value(result_dict, nodes_new_dict, link_dict): - source, source_handle = link_dict["source"], link_dict["sourcePort"] + source, source_handle = link_dict[SOURCE_LABEL], link_dict[SOURCE_PORT_LABEL] if source in result_dict.keys(): result = result_dict[source] elif source in nodes_new_dict.keys(): @@ -53,9 +65,9 @@ def load_workflow_json(file_name): with open(file_name, "r") as f: content = json.load(f) - edges_new_lst = content["edges"] + edges_new_lst = content[EDGES_LABEL] nodes_new_dict = {} - for k, v in convert_nodes_list_to_dict(nodes_list=content["nodes"]).items(): + for k, v in convert_nodes_list_to_dict(nodes_list=content[NODES_LABEL]).items(): if isinstance(v, str) and "." in v: p, m = v.rsplit('.', 1) mod = import_module(p) diff --git a/python_workflow_definition/src/python_workflow_definition/pyiron_base.py b/python_workflow_definition/src/python_workflow_definition/pyiron_base.py index 3e7b14e..b763d5e 100644 --- a/python_workflow_definition/src/python_workflow_definition/pyiron_base.py +++ b/python_workflow_definition/src/python_workflow_definition/pyiron_base.py @@ -6,7 +6,17 @@ from pyiron_base import job from pyiron_base.project.delayed import DelayedObject -from python_workflow_definition.shared import get_kwargs, get_source_handles, convert_nodes_list_to_dict +from python_workflow_definition.shared import ( + get_kwargs, + get_source_handles, + convert_nodes_list_to_dict, + NODES_LABEL, + EDGES_LABEL, + SOURCE_LABEL, + SOURCE_PORT_LABEL, + TARGET_LABEL, + TARGET_PORT_LABEL, +) def _resort_total_lst(total_lst, nodes_dict): @@ -16,7 +26,7 @@ def _resort_total_lst(total_lst, nodes_dict): while len(total_new_lst) < len(total_lst): for ind, connect in total_lst: if ind not in ordered_lst: - source_lst = [sd["source"] for sd in connect.values()] + source_lst = [sd[SOURCE_LABEL] for sd in connect.values()] if all([s in ordered_lst or s in nodes_without_dep_lst for s in source_lst]): ordered_lst.append(ind) total_new_lst.append([ind, connect]) @@ -24,15 +34,15 @@ def _resort_total_lst(total_lst, nodes_dict): def _group_edges(edges_lst): - edges_sorted_lst = sorted(edges_lst, key=lambda x: x["target"], reverse=True) + edges_sorted_lst = sorted(edges_lst, key=lambda x: x[TARGET_LABEL], reverse=True) total_lst, tmp_lst = [], [] - target_id = edges_sorted_lst[0]["target"] + target_id = edges_sorted_lst[0][TARGET_LABEL] for ed in edges_sorted_lst: - if target_id == ed["target"]: + if target_id == ed[TARGET_LABEL]: tmp_lst.append(ed) else: total_lst.append((target_id, get_kwargs(lst=tmp_lst))) - target_id = ed["target"] + target_id = ed[TARGET_LABEL] tmp_lst = [ed] total_lst.append((target_id, get_kwargs(lst=tmp_lst))) return total_lst @@ -54,8 +64,8 @@ def _get_delayed_object_dict(total_lst, nodes_dict, source_handle_dict, pyiron_p k: _get_source( nodes_dict=nodes_dict, delayed_object_dict=delayed_object_dict, - source=v["source"], - sourceHandle=v["sourcePort"], + source=v[SOURCE_LABEL], + sourceHandle=v[SOURCE_PORT_LABEL], ) for k, v in input_dict.items() } @@ -152,24 +162,24 @@ def _get_edges_dict(edges_lst, nodes_dict, connection_dict, lookup_dict): if isinstance(output, DelayedObject): if output._list_index is not None: edges_dict_lst.append({ - "target": target, - "targetPort": target_handle, - "source": connection_dict[output_name], - "sourcePort": f"s_{output._list_index}", # check for list index + TARGET_LABEL: target, + TARGET_PORT_LABEL: target_handle, + SOURCE_LABEL: connection_dict[output_name], + SOURCE_PORT_LABEL: f"s_{output._list_index}", # check for list index }) else: edges_dict_lst.append({ - "target": target, - "targetPort": target_handle, - "source": connection_dict[output_name], - "sourcePort": output._output_key, # check for list index + TARGET_LABEL: target, + TARGET_PORT_LABEL: target_handle, + SOURCE_LABEL: connection_dict[output_name], + SOURCE_PORT_LABEL: output._output_key, # check for list index }) else: edges_dict_lst.append({ - "target": target, - "targetPort": target_handle, - "source": connection_dict[output_name], - "sourcePort": None, + TARGET_LABEL: target, + TARGET_PORT_LABEL: target_handle, + SOURCE_LABEL: connection_dict[output_name], + SOURCE_PORT_LABEL: None, }) existing_connection_lst.append(connection_name) return edges_dict_lst @@ -179,9 +189,9 @@ def load_workflow_json(project, file_name): with open(file_name, "r") as f: content = json.load(f) - edges_new_lst = content["edges"] + edges_new_lst = content[EDGES_LABEL] nodes_new_dict = {} - for k, v in convert_nodes_list_to_dict(nodes_list=content["nodes"]).items(): + for k, v in convert_nodes_list_to_dict(nodes_list=content[NODES_LABEL]).items(): if isinstance(v, str) and "." in v: p, m = v.rsplit('.', 1) if p == "python_workflow_definition.shared": @@ -224,4 +234,4 @@ def write_workflow_json(delayed_object, file_name="workflow.json"): nodes_store_lst.append({"id": k, "value": v}) with open(file_name, "w") as f: - json.dump({"nodes": nodes_store_lst, "edges": edges_new_lst}, f) + json.dump({NODES_LABEL: nodes_store_lst, EDGES_LABEL: edges_new_lst}, f) diff --git a/python_workflow_definition/src/python_workflow_definition/shared.py b/python_workflow_definition/src/python_workflow_definition/shared.py index ca7cbe8..f2f71d5 100644 --- a/python_workflow_definition/src/python_workflow_definition/shared.py +++ b/python_workflow_definition/src/python_workflow_definition/shared.py @@ -1,3 +1,11 @@ +NODES_LABEL = "nodes" +EDGES_LABEL = "edges" +SOURCE_LABEL = "source" +SOURCE_PORT_LABEL = "sourcePort" +TARGET_LABEL = "target" +TARGET_PORT_LABEL = "targetPort" + + def get_dict(**kwargs): # NOTE: In WG, this will automatically be wrapped in a dict with the `result` key return {k: v for k, v in kwargs.items()} @@ -9,16 +17,20 @@ def get_list(**kwargs): def get_kwargs(lst): - return {t["targetPort"]: {"source": t["source"], "sourcePort": t["sourcePort"]} for t in lst} + return { + t[TARGET_PORT_LABEL]: { + SOURCE_LABEL: t[SOURCE_LABEL], + SOURCE_PORT_LABEL: t[SOURCE_PORT_LABEL] + } for t in lst + } def get_source_handles(edges_lst): source_handle_dict = {} for ed in edges_lst: - if ed["source"] not in source_handle_dict.keys(): - source_handle_dict[ed["source"]] = [ed["sourcePort"]] - else: - source_handle_dict[ed["source"]].append(ed["sourcePort"]) + if ed[SOURCE_LABEL] not in source_handle_dict.keys(): + source_handle_dict[ed[SOURCE_LABEL]] = [] + source_handle_dict[ed[SOURCE_LABEL]].append(ed[SOURCE_PORT_LABEL]) return { k: list(range(len(v))) if len(v) > 1 and all([el is None for el in v]) else v for k, v in source_handle_dict.items()