flame: returning all batch nodes

This commit is contained in:
Jakub Jezek 2022-04-07 15:05:46 +02:00
parent cd59b3af66
commit 4ed6c0257e
No known key found for this signature in database
GPG key ID: D8548FBF690B100A

View file

@ -91,6 +91,9 @@ def create_batch_group_conent(batch_nodes, batch_links, batch_group=None):
batch_nodes (list of dict): each dict is node definition batch_nodes (list of dict): each dict is node definition
batch_links (list of dict): each dict is link definition batch_links (list of dict): each dict is link definition
batch_group (PyBatch, optional): batch group. Defaults to None. batch_group (PyBatch, optional): batch group. Defaults to None.
Return:
dict: all batch nodes {name or id: PyNode}
""" """
# make sure some batch obj is present # make sure some batch obj is present
batch_group = batch_group or flame.batch batch_group = batch_group or flame.batch
@ -98,7 +101,6 @@ def create_batch_group_conent(batch_nodes, batch_links, batch_group=None):
b.name.get_value(): b b.name.get_value(): b
for b in batch_group.nodes for b in batch_group.nodes
} }
created_nodes = {}
for node in batch_nodes: for node in batch_nodes:
# NOTE: node_props needs to be ideally OrederDict type # NOTE: node_props needs to be ideally OrederDict type
node_id, node_type, node_props = ( node_id, node_type, node_props = (
@ -121,7 +123,7 @@ def create_batch_group_conent(batch_nodes, batch_links, batch_group=None):
setattr(batch_node, key, value) setattr(batch_node, key, value)
# add created node for possible linking # add created node for possible linking
created_nodes[node_id] = batch_node all_batch_nodes[node_id] = batch_node
# link nodes to each other # link nodes to each other
for link in batch_links: for link in batch_links:
@ -129,16 +131,18 @@ def create_batch_group_conent(batch_nodes, batch_links, batch_group=None):
# check if all linking nodes are available # check if all linking nodes are available
if not all([ if not all([
created_nodes.get(_from_n["id"]), all_batch_nodes.get(_from_n["id"]),
created_nodes.get(_to_n["id"]) all_batch_nodes.get(_to_n["id"])
]): ]):
continue continue
# link nodes in defined link # link nodes in defined link
batch_group.connect_nodes( batch_group.connect_nodes(
created_nodes[_from_n["id"]], _from_n["connector"], all_batch_nodes[_from_n["id"]], _from_n["connector"],
created_nodes[_to_n["id"]], _to_n["connector"] all_batch_nodes[_to_n["id"]], _to_n["connector"]
) )
# sort batch nodes # sort batch nodes
batch_group.organize() batch_group.organize()
return all_batch_nodes