From 4ed6c0257ece24cc254afb4edc9a4cc698c4dd53 Mon Sep 17 00:00:00 2001
From: Jakub Jezek <jakub@orbi.tools>
Date: Thu, 7 Apr 2022 15:05:46 +0200
Subject: [PATCH] flame: returning all batch nodes

---
 openpype/hosts/flame/api/batch_utils.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/openpype/hosts/flame/api/batch_utils.py b/openpype/hosts/flame/api/batch_utils.py
index 23d16f8d1a..20118c249c 100644
--- a/openpype/hosts/flame/api/batch_utils.py
+++ b/openpype/hosts/flame/api/batch_utils.py
@@ -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_links (list of dict): each dict is link definition
         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
     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
         for b in batch_group.nodes
     }
-    created_nodes = {}
     for node in batch_nodes:
         # NOTE: node_props needs to be ideally OrederDict type
         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)
 
         # add created node for possible linking
-        created_nodes[node_id] = batch_node
+        all_batch_nodes[node_id] = batch_node
 
     # link nodes to each other
     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
         if not all([
-            created_nodes.get(_from_n["id"]),
-            created_nodes.get(_to_n["id"])
+            all_batch_nodes.get(_from_n["id"]),
+            all_batch_nodes.get(_to_n["id"])
         ]):
             continue
 
         # link nodes in defined link
         batch_group.connect_nodes(
-            created_nodes[_from_n["id"]], _from_n["connector"],
-            created_nodes[_to_n["id"]], _to_n["connector"]
+            all_batch_nodes[_from_n["id"]], _from_n["connector"],
+            all_batch_nodes[_to_n["id"]], _to_n["connector"]
         )
 
     # sort batch nodes
     batch_group.organize()
+
+    return all_batch_nodes