From e3f1a446689d25ceaa2500a11386e79ae2823de3 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 1 Mar 2023 22:20:56 -0500 Subject: [PATCH] Handle lists of arguments (like 'Join Geometry') --- operators/convert_tree.py | 52 +++++++++++++++++++++++++++++---------- 1 file changed, 39 insertions(+), 13 deletions(-) diff --git a/operators/convert_tree.py b/operators/convert_tree.py index 7090a29..ea30215 100644 --- a/operators/convert_tree.py +++ b/operators/convert_tree.py @@ -15,8 +15,20 @@ class _Assignment: self.props = props self.arguments = {} self.argument_dot_access = {} + + def _flattened_arguments(self): + for a in self.arguments.values(): + if isinstance(a, _Assignment): + yield a + elif isinstance(a, list): + yield from a + def flattened_arguments(self): + return [*self._flattened_arguments()] - def convert_argument(self, k, v, delimiter='='): + def convert_argument(self, k, v, delimiter='=', index=None): + if isinstance(v, list): + v = ', '.join([self.convert_argument("", sv, delimiter="", index=i) for i, sv in enumerate(v)]) + return f"{k}{delimiter}[{v}]" if not isinstance(v, _Assignment): if isinstance(v, str): v = f'"{v}"' @@ -24,7 +36,7 @@ class _Assignment: v = str(v) return f"{k}{delimiter}{v}" if v.node.type == 'GROUP_INPUT': - return f"{k}{delimiter}{self.argument_dot_access[k]}" + return f"{k}{delimiter}{self.argument_dot_access[k] if index is None else self.argument_dot_access[k][index]}" else: return f"{k}{delimiter}{v.name}{'.' + self.argument_dot_access[k] if len(list(o for o in v.node.outputs if o.enabled)) > 1 else ''}" @@ -32,6 +44,7 @@ class _Assignment: snake_case_name = self.node.bl_rna.name.lower().replace(' ', '_') args = ', '.join([ self.convert_argument(k, v) + for k, v in (list(self.props.items()) + list(self.arguments.items())) ]) return f"{self.name} = {snake_case_name}({args})" @@ -71,8 +84,20 @@ class ConvertTree(bpy.types.Operator): output_name = link.from_socket.name.lower().replace(' ', '_') if argument_name in to_node.props: del to_node.props[argument_name] - to_node.arguments[argument_name] = from_node - to_node.argument_dot_access[argument_name] = output_name + if argument_name in to_node.arguments: + if isinstance(to_node.arguments[argument_name], list): + index = len(to_node.arguments[argument_name]) + to_node.arguments[argument_name].append(from_node) + to_node.argument_dot_access[argument_name][index] = output_name + else: + to_node.arguments[argument_name] = [to_node.arguments[argument_name], from_node] + to_node.argument_dot_access[argument_name] = { + 0: to_node.argument_dot_access[argument_name], + 1: output_name + } + else: + to_node.arguments[argument_name] = from_node + to_node.argument_dot_access[argument_name] = output_name def topological_sort(root): seen = set() @@ -83,9 +108,9 @@ class ConvertTree(bpy.types.Operator): v = q.pop() if v not in seen: seen.add(v) - q.extend([a for a in v.arguments.values() if isinstance(a, _Assignment)]) + q.extend(v.flattened_arguments()) - while stack and v not in stack[-1].arguments.values(): + while stack and v not in stack[-1].flattened_arguments(): order.append(stack.pop()) stack.append(v) @@ -96,14 +121,15 @@ class ConvertTree(bpy.types.Operator): body = '\n '.join([a.to_script() for a in sorted_assignments]) - group_input = next(n for n in tree.nodes if n.type == 'GROUP_INPUT') + group_input = next((n for n in tree.nodes if n.type == 'GROUP_INPUT'), None) tree_arguments = [] - for output in group_input.outputs: - if output.type == 'CUSTOM': - continue - output_name = output.name.lower().replace(' ', '_') - output_type = output.bl_rna.identifier.replace('NodeSocket', '') - tree_arguments.append(f"{output_name}: {output_type}") + if group_input is not None: + for output in group_input.outputs: + if output.type == 'CUSTOM': + continue + output_name = output.name.lower().replace(' ', '_') + output_type = output.bl_rna.identifier.replace('NodeSocket', '') + tree_arguments.append(f"{output_name}: {output_type}") tree_arguments = ', '.join(tree_arguments) tree_returns = []