from .utilities import (prRed, prGreen, prPurple, prWhite, prOrange, wrapRed, wrapGreen, wrapPurple, wrapWhite, wrapOrange,) from .utilities import init_connections, init_dependencies from .utilities import class_for_mantis_prototype_node from .base_definitions import SchemaUINode, GraphError, replace_types, custom_props_types from .node_container_common import setup_custom_props_from_np # a class that solves Schema nodes class SchemaSolver: def __init__(self, schema_dummy, nodes, prototype, signature=None): self.all_nodes = nodes # this is the parsed tree from Mantis self.node = schema_dummy self.tree = prototype.node_tree self.uuid = self.node.uuid if signature: self.signature = signature else: self.signature = self.node.signature self.schema_nodes={} self.solved_nodes = {} self.incoming_connections = {} self.outgoing_connections = {} self.constant_in = {} self.constant_out = {} self.array_input_connections = {} self.array_output_connections = {} self.nested_schemas = {} self.autogenerated_nodes = {} self.held_links = [] self.tree_path_names = [*self.node.signature] # same tree as the schema node self.autogen_path_names = ['SCHEMA_AUTOGENERATED', *self.node.signature[1:]] self.index_link = self.node.inputs['Schema Length'].links[0] self.solve_length = self.node.evaluate_input("Schema Length") # I'm making this a property of the solver because the solver's data is modified as it solves each iteration # So self.index = 0 self.init_schema_links() self.set_index_strings() for ui_node in self.tree.nodes: if isinstance(ui_node, SchemaUINode): # first we need to fill the parameters of the schema nodes. # we use the bl_idname because all schema nodes should be single-instance signature = (*self.tree_path_names, ui_node.bl_idname) # We use the solver's signature here because it represents the "original" signature of the schema UI group node # since this schema solver may be in a nested schema, and its node's signature may have uuid/index attached. get_sig = (*self.signature, ui_node.bl_idname) if not (mantis_node := self.all_nodes.get(get_sig)): raise RuntimeError(wrapRed(f"Not found: {get_sig}")) self.schema_nodes[signature] = mantis_node mantis_node.fill_parameters(ui_node) def set_index_strings(self): self.index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index).zfill(4) self.prev_index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index-1).zfill(4) def init_schema_links(self,): """ Sort and store the links to/from the Schema group node.""" for item in self.tree.interface.items_tree: if item.item_type == 'PANEL': continue if not item.parent: raise GraphError("ERROR: Schema tree has inputs that are not in categories. This is not supported") match item.parent.name: case 'Connection': if item.in_out == 'INPUT': if incoming_links := self.node.inputs[item.identifier].links: self.incoming_connections[item.name] = incoming_links[0] else: self.incoming_connections[item.name] = None else: # OUTPUT if outgoing_links := self.node.outputs[item.identifier].links: self.outgoing_connections[item.name] = outgoing_links.copy() else: self.outgoing_connections[item.name] = [] case 'Constant': if item.in_out == 'INPUT': if constant_in_links := self.node.inputs[item.identifier].links: self.constant_in[item.name] = constant_in_links[0] else: self.constant_in[item.name] = None else: # OUTPUT if constant_out_links := self.node.outputs[item.identifier].links: self.constant_out[item.name] = constant_out_links.copy() else: self.constant_out[item.name] = [] case 'Array': if item.in_out == 'INPUT': if item.identifier not in self.array_input_connections.keys(): self.array_input_connections[item.identifier]=[] if in_links := self.node.inputs[item.identifier].links: self.array_input_connections[item.identifier]=in_links.copy() else: # OUTPUT if item.identifier not in self.array_output_connections.keys(): self.array_output_connections[item.identifier]=[] if out_links := self.node.outputs[item.identifier].links: self.array_output_connections[item.identifier] = out_links.copy() def gen_solve_iteration_mantis_nodes(self, frame_mantis_nodes, unprepared): for prototype_ui_node in self.tree.nodes: if isinstance(prototype_ui_node, SchemaUINode): continue # IGNORE the schema interface nodes, we already made them in __init__() # they are reused for each iteration. elif prototype_ui_node.bl_idname in ['NodeFrame', 'NodeReroute']: continue # IGNORE stuff that is purely UI - frames, reroutes. signature = (*self.autogen_path_names, prototype_ui_node.name+self.index_str()) prototype_mantis_node = self.all_nodes[(*self.signature, prototype_ui_node.name)] # the prototype_mantis_node was generated inside the schema when we parsed the tree. # it is the prototype of the mantis node which we make for this iteration # for Schema sub-nodes ... they need a prototype to init. if prototype_mantis_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']: from .utilities import get_node_prototype ui_node = get_node_prototype(prototype_mantis_node.signature, prototype_mantis_node.base_tree) if prototype_mantis_node.node_type == 'DUMMY_SCHEMA': mantis_node = prototype_mantis_node.__class__( signature, prototype_mantis_node.base_tree, prototype=ui_node, natural_signature = (*self.node.signature, ui_node.name) ) else: mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree, prototype=ui_node) else: mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree) frame_mantis_nodes[mantis_node.signature] = mantis_node if mantis_node.prepared == False: unprepared.append(mantis_node) if mantis_node.__class__.__name__ in custom_props_types: setup_custom_props_from_np(mantis_node, prototype_ui_node) mantis_node.fill_parameters(prototype_ui_node) def handle_link_from_index_input(self, index, frame_mantis_nodes, ui_link): from .utilities import get_link_in_out _from_name, to_name = get_link_in_out(ui_link) to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ] if to_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']: from .utilities import gen_nc_input_for_data nc_cls = gen_nc_input_for_data(ui_link.from_socket) if (nc_cls): #HACK sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str(), ui_link.from_socket.name, ui_link.from_socket.identifier) nc_from = nc_cls(sig, self.node.base_tree) # ugly! maybe even a HACK! nc_from.inputs = {} from .base_definitions import NodeSocket nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)} from .base_definitions import get_socket_value nc_from.parameters = {ui_link.from_socket.name:index} frame_mantis_nodes[sig]=nc_from from_node = nc_from self.solved_nodes[sig]=from_node _connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=ui_link.to_socket.identifier) return # Since the index is already determined, it is safe to remove the socket and just keep the value. to_node.parameters[ui_link.to_socket.name] = index del to_node.inputs[ui_link.to_socket.name] def handle_link_from_schema_length_input(self, frame_mantis_nodes, ui_link): from .utilities import get_link_in_out # see, here I can just use the schema node _from_name, to_name = get_link_in_out(ui_link) to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())] # this self.index_link is only used here? if (self.index_link.from_node): connection = self.index_link.from_node.outputs[self.index_link.from_socket].connect(node=to_node, socket=ui_link.to_socket.name) # otherwise we can autogen an int input I guess...? else: raise RuntimeError("I was expecting there to be an incoming connection here for Schema Length") def handle_link_from_incoming_connection_input(self, frame_mantis_nodes, ui_link): from .utilities import get_link_in_out incoming = self.incoming_connections[ui_link.from_socket.name] from_node = incoming.from_node _from_name, to_name = get_link_in_out(ui_link) to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ] connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name) init_connections(from_node) def handle_link_from_constant_input(self, frame_mantis_nodes, ui_link, to_ui_node): from .utilities import get_link_in_out incoming = self.constant_in[ui_link.from_socket.name] from_node = incoming.from_node to_name = get_link_in_out(ui_link)[1] to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())] to_socket=ui_link.to_socket.name from .base_definitions import SchemaGroup if isinstance(to_ui_node, SchemaGroup): to_socket=ui_link.to_socket.identifier connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=to_socket) init_connections(from_node) def handle_link_to_array_input_get(self, frame_mantis_nodes, ui_link, index): from .utilities import get_link_in_out from_name, to_name = get_link_in_out(ui_link) from_nc = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str())] to_nc = self.schema_nodes[(*self.tree_path_names, to_name)] # this only needs to be done once: if index == 0: # BUG? HACK? TODO find out what is going on here. # Kill the link between the schema node group and the node connecting to it old_nc = self.all_nodes[(*self.tree_path_names, from_name)] # I am not sure about this! existing_link = old_nc.outputs[ui_link.from_socket.name].links[0] existing_link.die() # connection = from_nc.outputs[ui_link.from_socket.name].connect(node=to_nc, socket=ui_link.to_socket.name) def handle_link_from_array_input(self, frame_mantis_nodes, ui_link, index): from .utilities import get_link_in_out get_index = index try: incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index] except IndexError: if len(self.array_input_connections[ui_link.from_socket.identifier]) > 0: incoming = self.array_input_connections[ui_link.from_socket.identifier][0] # prOrange(incoming.from_node.node_type) if incoming.from_node.node_type not in ['DUMMY_SCHEMA']: raise NotImplementedError(wrapRed("dev: make it so Mantis checks if there are enough Array inputs.")) else: # do nothing return else: raise RuntimeError(wrapRed("make it so Mantis checks if there are enough Array inputs!")) to_name = get_link_in_out(ui_link)[1] to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())] connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name) init_connections(incoming.from_node) def handle_link_to_constant_output(self, frame_mantis_nodes, index, ui_link, to_ui_node): from .utilities import get_link_in_out to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)] expose_when = to_node.evaluate_input('Expose when N==') if index == expose_when: for outgoing in self.constant_out[ui_link.to_socket.name]: to_node = outgoing.to_node from_name = get_link_in_out(ui_link)[0] from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str()) ] connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=outgoing.to_socket) # WTF is even happening here?? TODO BUG HACK def handle_link_to_array_output(self, frame_mantis_nodes, index, ui_link, to_ui_node, from_ui_node):# if this duplicated code works, dedupe! from .utilities import get_link_in_out to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)] # get it by [], we want a KeyError if this fails for outgoing in self.array_output_connections[ui_link.to_socket.identifier]: # print (type(outgoing)) from .schema_containers import SchemaIndex from_name = get_link_in_out(ui_link)[0] from_node = frame_mantis_nodes[ (*self.autogen_path_names, from_name+self.index_str()) ] if not from_node: from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)] to_node = outgoing.to_node if isinstance(from_node, SchemaIndex): # I think I need to dedup this stuff # print("INDEX") from .utilities import gen_nc_input_for_data nc_cls = gen_nc_input_for_data(ui_link.from_socket) if (nc_cls): #HACK sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str(), ui_link.from_socket.name, ui_link.from_socket.identifier) nc_from = nc_cls(sig, self.node.base_tree) # ugly! maybe even a HACK! nc_from.inputs = {} from .node_container_common import NodeSocket nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)} from .node_container_common import get_socket_value if ui_link.from_socket.name in ['Index']: nc_from.parameters = {ui_link.from_socket.name:index} else: nc_from.parameters = {ui_link.from_socket.name:self.solve_length} frame_mantis_nodes[sig]=nc_from from_node = nc_from self.solved_nodes[sig]=from_node # I have a feeling that something bad will happen if both of these conditions (above and below) are true if to_node.node_type == 'DUMMY_SCHEMA' and to_node.prepared: other_stem = ('SCHEMA_AUTOGENERATED', *to_node.signature[1:]) from .utilities import get_node_prototype other_schema_np = get_node_prototype(to_node.signature, to_node.base_tree) other_schema_tree = other_schema_np.node_tree for n in other_schema_tree.nodes: if n.bl_idname not in ["SchemaArrayInput", "SchemaArrayInputGet"]: continue out = n.outputs[outgoing.to_socket] for l in out.links: other_index_str = lambda : '.'+str(to_node.uuid)+'.'+str(index).zfill(4) # get it by [], we want a KeyError if this fails try: out_node = self.all_nodes[(*other_stem, l.to_node.name+other_index_str())] except KeyError as e: for n in self.all_nodes: if len(n) > len(other_stem)+1: break for elem in other_stem: if elem not in n: break else: print(n) raise e connection = from_node.outputs[ui_link.from_socket.name].connect(node=out_node, socket=l.to_socket.name) else: connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=outgoing.to_socket) def handle_link_from_array_input_get(self, frame_mantis_nodes, index, ui_link, from_ui_node ): from .utilities import get_link_in_out get_index = index from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)] from .utilities import cap, wrap get_index = from_node.evaluate_input("Index", index) oob = from_node.evaluate_input("OoB Behaviour") # we must assume that the array has sent the correct number of links if oob == 'WRAP': get_index = wrap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1, 0) if oob == 'HOLD': get_index = cap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1) try: incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index] except IndexError: raise RuntimeError(wrapRed("Dummy! You need to make it so Mantis checks if there are enough Array inputs! It should probably have a Get Index!")) to_name = get_link_in_out(ui_link)[1] to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())] connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name) init_connections(incoming.from_node) def solve_iteration(self): """ Solve an iteration of the schema. - 1 Create the Mantis Node instances that represent this iteration of the schema - 2 Connect the links from the entrypoint or previous iteration. - 3 Connect the constant and array links, and any link between nodes entirely within the tree - 4 Prepare the nodes that modify data (in case of e.g. array get index or nested schema length input) - 5 Connect the final prepared nodes and return the nodes that were created in this schema iteration (frame). This function also adds to held_links to pass data between iterations. """ from .schema_definitions import (SchemaIndex, SchemaArrayInput, SchemaArrayInputGet, SchemaArrayOutput, SchemaConstInput, SchemaConstOutput, SchemaOutgoingConnection, SchemaIncomingConnection,) from .utilities import clear_reroutes from .utilities import get_link_in_out, link_node_containers self.set_index_strings() frame_mantis_nodes = {} # Later we have to run bPrepare() on these guys, so we make the deque and fill it now. from collections import deque unprepared= deque() self.gen_solve_iteration_mantis_nodes(frame_mantis_nodes, unprepared) # This is where we handle node connections BETWEEN frames while(self.held_links): ui_link = self.held_links.pop() to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node if isinstance(to_ui_node, SchemaOutgoingConnection): mantis_incoming_node = self.schema_nodes[*self.tree_path_names, 'SchemaIncomingConnection'] for mantis_link in mantis_incoming_node.outputs[ui_link.to_socket.name].links: to_mantis_node, to_mantis_socket = mantis_link.to_node, mantis_link.to_socket from_name = get_link_in_out(ui_link)[0] from_mantis_node = self.solved_nodes[ (*self.autogen_path_names, from_name+self.prev_index_str()) ] to_mantis_node = frame_mantis_nodes[ (*self.autogen_path_names, to_mantis_node.signature[-1]+self.index_str()) ] connection = from_mantis_node.outputs[ui_link.from_socket.name].connect(node=to_mantis_node, socket=to_mantis_socket) # We want to delete the links from the tree into the schema node. # TODO: this is not robust enough and I do not feel sure this is doing the right thing. if existing_link := self.incoming_connections[ui_link.to_socket.name]: if existing_link.to_node == self.node: print ("Deleting...", existing_link) if self.node.signature[-1] in existing_link.to_node.signature: existing_link.die() # BUG may exist here. self.incoming_connections[ui_link.to_socket.name] = connection # Get the rerouted links from the graph. We don't really need to do this every iteration. # TODO: use profiling to determine if this is slow; if so: copy & reuse the data, refactor the pop()'s out. ui_links = clear_reroutes(list(self.tree.links)) # Now we handle ui_links in the current frame, including those ui_links between Schema nodes and "real" nodes awaiting_prep_stage = [] for ui_link in ui_links: to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node if isinstance(from_ui_node, SchemaIndex): if ui_link.from_socket.name == "Index": self.handle_link_from_index_input(self.index, frame_mantis_nodes, ui_link) elif ui_link.from_socket.name == "Schema Length": self.handle_link_from_schema_length_input(frame_mantis_nodes, ui_link) continue if isinstance(from_ui_node, SchemaIncomingConnection): if ui_link.from_socket.name in self.incoming_connections.keys(): self.handle_link_from_incoming_connection_input(frame_mantis_nodes, ui_link) continue if isinstance(from_ui_node, SchemaConstInput): if ui_link.from_socket.name in self.constant_in.keys(): self.handle_link_from_constant_input( frame_mantis_nodes, ui_link, to_ui_node) continue if isinstance(to_ui_node, SchemaArrayInputGet): self.handle_link_to_array_input_get( frame_mantis_nodes, ui_link, self.index) continue if isinstance(from_ui_node, SchemaArrayInput): self.handle_link_from_array_input(frame_mantis_nodes, ui_link, self.index) continue # HOLD these links to the next iteration: if isinstance(to_ui_node, SchemaOutgoingConnection): self.held_links.append(ui_link) continue # HOLD these links until prep is done a little later if isinstance(to_ui_node, SchemaConstOutput) or isinstance(to_ui_node, SchemaArrayOutput) or \ isinstance(from_ui_node, SchemaArrayInputGet): awaiting_prep_stage.append(ui_link) continue # for any of the special cases, we hit a 'continue' block. So this connection is not special, and is made here. connection = link_node_containers(self.autogen_path_names, ui_link, frame_mantis_nodes, from_suffix=self.index_str(), to_suffix=self.index_str()) for k,v in frame_mantis_nodes.items(): self.solved_nodes[k]=v init_dependencies(v) # it is hard to overstate how important this single line of code is # This while-loop is a little scary, but in my testing it has never been a problem. # At this point, we've already run a pretty exhaustive preperation phase to prep the schema's dependencies # So at this point, if this while loop hangs it is because of an error elsewhere. while unprepared: nc = unprepared.pop() if sum([dep.prepared for dep in nc.hierarchy_dependencies]) == len(nc.hierarchy_dependencies): nc.bPrepare() if nc.node_type == 'DUMMY_SCHEMA': self.solve_nested_schema(nc) else: unprepared.appendleft(nc) # just rotate them until they are ready. while(awaiting_prep_stage): ui_link = awaiting_prep_stage.pop() to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node if isinstance(to_ui_node, SchemaConstOutput): self.handle_link_to_constant_output(frame_mantis_nodes, self.index, ui_link, to_ui_node) if isinstance(to_ui_node, SchemaArrayOutput): self.handle_link_to_array_output(frame_mantis_nodes, self.index, ui_link, to_ui_node, from_ui_node) if isinstance(from_ui_node, SchemaArrayInputGet): self.handle_link_from_array_input_get(frame_mantis_nodes, self.index, ui_link, from_ui_node ) # end seciton return frame_mantis_nodes def solve_nested_schema(self, schema_nc): """ Solves all schema node groups found in this Schema. This is a recursive function, which will solve all levels of nested schema - since this function is called by solver.solve(). """ if schema_nc.prepared == False: all_nodes = self.all_nodes.copy() ui_node = schema_nc.prototype length = schema_nc.evaluate_input("Schema Length") tree = ui_node.node_tree prOrange(f"Expanding schema {tree.name} in node {schema_nc} with length {length}.") solver = SchemaSolver(schema_nc, all_nodes, ui_node, schema_nc.natural_signature) solved_nodes = solver.solve() schema_nc.prepared = True for k,v in solved_nodes.items(): self.solved_nodes[k]=v def finalize(self, frame_nc): from .schema_definitions import (SchemaOutgoingConnection,) for i in range(len(self.held_links)): link = self.held_links.pop() to_np = link.to_socket.node; from_np = link.from_socket.node if isinstance(to_np, SchemaOutgoingConnection): if link.to_socket.name in self.outgoing_connections.keys(): if (outgoing_links := self.outgoing_connections[link.to_socket.name]) is None: continue for outgoing in outgoing_links: if outgoing: to_node = outgoing.to_node from_node =frame_nc[(*self.autogen_path_names, from_np.name+self.index_str()) ] connection = from_node.outputs[link.from_socket.name].connect(node=to_node, socket=outgoing.to_socket) # we need to kill the link between the Schema itself and the next node and update the deps. Otherwise:confusing bugs. outgoing.die(); init_dependencies(to_node) # else: # the node just isn't connected out this socket. # # solve all unsolved nested schemas... for schema_sig, schema_nc in self.nested_schemas.items(): self.solve_nested_schema(schema_nc) for n in self.autogenerated_nodes.values(): init_connections(n) for c in n.connections: init_dependencies(c) all_outgoing_links = [] for conn in self.outgoing_connections.values(): for outgoing in conn: all_outgoing_links.append(outgoing) for conn in self.constant_out.values(): for outgoing in conn: all_outgoing_links.append(outgoing) for conn in self.array_output_connections.values(): for outgoing in conn: all_outgoing_links.append(outgoing) for outgoing in all_outgoing_links: to_node = outgoing.to_node for l in to_node.inputs[outgoing.to_socket].links: if self.node == l.from_node: l.die() for inp in self.node.inputs.values(): for l in inp.links: init_connections(l.from_node) # to force it to have hierarchy connections with the new nodes. def solve(self): for index in range(self.solve_length): self.index = index frame_mantis_nodes = self.solve_iteration() for sig, nc in frame_mantis_nodes.items(): if nc.node_type == 'DUMMY_SCHEMA': self.nested_schemas[sig] = nc self.finalize(frame_mantis_nodes) self.node.solver = self return self.solved_nodes