base_definitions.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693
  1. #Mantis Nodes Base
  2. import bpy
  3. from bpy.props import (BoolProperty, StringProperty, EnumProperty, CollectionProperty, \
  4. IntProperty, IntVectorProperty, PointerProperty, BoolVectorProperty)
  5. from . import ops_nodegroup
  6. from bpy.types import NodeTree, Node, PropertyGroup, Operator, UIList, Panel
  7. from .utilities import (prRed, prGreen, prPurple, prWhite,
  8. prOrange,
  9. wrapRed, wrapGreen, wrapPurple, wrapWhite,
  10. wrapOrange,)
  11. from .utilities import get_socket_maps, relink_socket_map, do_relink
  12. def TellClasses():
  13. #Why use a function to do this? Because I don't need every class to register.
  14. return [ MantisTree,
  15. SchemaTree,
  16. MantisNodeGroup,
  17. SchemaGroup,
  18. ]
  19. def error_popup_draw(self, context):
  20. self.layout.label(text="Error executing tree. See Console.")
  21. mantis_root = ".".join(__name__.split('.')[:-1]) # absolute HACK
  22. class MantisTree(NodeTree):
  23. '''A custom node tree type that will show up in the editor type list'''
  24. bl_idname = 'MantisTree'
  25. bl_label = "Rigging Nodes"
  26. bl_icon = 'OUTLINER_OB_ARMATURE'
  27. tree_valid:BoolProperty(default=False)
  28. do_live_update:BoolProperty(default=True) # use this to disable updates for e.g. scripts
  29. num_links:IntProperty(default=-1)
  30. filepath:StringProperty(default="", subtype='FILE_PATH')
  31. is_executing:BoolProperty(default=False)
  32. is_exporting:BoolProperty(default=False)
  33. execution_id:StringProperty(default='')
  34. mantis_version:IntVectorProperty(default=[0,9,2])
  35. # this prevents the node group from executing on the next depsgraph update
  36. # because I don't always have control over when the dg update happens.
  37. prevent_next_exec:BoolProperty(default=False)
  38. parsed_tree={}
  39. if (bpy.app.version < (4, 4, 0)): # in 4.4 this leads to a crash
  40. @classmethod
  41. def valid_socket_type(cls : NodeTree, socket_type: str):
  42. # https://docs.blender.org/api/master/bpy.types.NodeTree.html#bpy.types.NodeTree.valid_socket_type
  43. from .socket_definitions import Tell_bl_idnames
  44. return socket_type in Tell_bl_idnames()
  45. # thank you, Sverchok
  46. def update_tree(self, context = None):
  47. if self.is_exporting:
  48. return
  49. # return
  50. self.is_executing = True
  51. from . import readtree
  52. prGreen("Validating Tree: %s" % self.name)
  53. try:
  54. self.parsed_tree = readtree.parse_tree(self)
  55. if context:
  56. self.display_update(context)
  57. self.is_executing = False
  58. self.tree_valid = True
  59. except GraphError as e:
  60. prRed("Failed to update node tree due to error.")
  61. self.tree_valid = False
  62. self.is_executing = False
  63. raise e
  64. finally:
  65. self.is_executing = False
  66. def display_update(self, context):
  67. if self.is_exporting:
  68. return
  69. self.is_executing = True
  70. current_tree = bpy.context.space_data.path[-1].node_tree
  71. for node in current_tree.nodes:
  72. if hasattr(node, "display_update"):
  73. try:
  74. node.display_update(self.parsed_tree, context)
  75. except Exception as e:
  76. print("Node \"%s\" failed to update display with error: %s" %(wrapGreen(node.name), wrapRed(e)))
  77. self.is_executing = False
  78. # TODO: deal with invalid links properly.
  79. # - Non-hierarchy links should be ignored in the circle-check and so the links should be marked valid in such a circle
  80. # - hierarchy-links should be marked invalid and prevent the tree from executing.
  81. def execute_tree(self,context, error_popups = False):
  82. self.prevent_next_exec = False
  83. if self.is_exporting:
  84. return
  85. # return
  86. prGreen("Executing Tree: %s" % self.name)
  87. self.is_executing = True
  88. from . import readtree
  89. try:
  90. readtree.execute_tree(self.parsed_tree, self, context, error_popups)
  91. except RecursionError as e:
  92. prRed("Recursion error while parsing tree.")
  93. finally:
  94. self.is_executing = False
  95. class SchemaTree(NodeTree):
  96. '''A node tree representing a schema to generate a Mantis tree'''
  97. bl_idname = 'SchemaTree'
  98. bl_label = "Rigging Nodes Schema"
  99. bl_icon = 'RIGID_BODY_CONSTRAINT'
  100. # these are only needed for consistent interface, but should not be used
  101. do_live_update:BoolProperty(default=True) # default to true so that updates work
  102. is_executing:BoolProperty(default=False)
  103. is_exporting:BoolProperty(default=False)
  104. mantis_version:IntVectorProperty(default=[0,9,2])
  105. if (bpy.app.version < (4, 4, 0)): # in 4.4 this leads to a crash
  106. @classmethod
  107. def valid_socket_type(cls, socket_type: str):
  108. # https://docs.blender.org/api/master/bpy.types.NodeTree.html#bpy.types.NodeTree.valid_socket_type
  109. from .socket_definitions import Tell_bl_idnames
  110. return socket_type in Tell_bl_idnames()
  111. # thank you, Sverchok
  112. #TODO: do a better job explaining how MantisNode and MantisUINode relate.
  113. class MantisUINode:
  114. """
  115. This class contains the common user-interface features of Mantis nodes.
  116. MantisUINode objects will spawn one or more MantisNode objects when the graph is evaluated.
  117. The MantisNode objects will pull the data from the UI node and use it to generate the graph.
  118. """
  119. mantis_node_library=''
  120. mantis_node_class_name=''
  121. mantis_class=None
  122. @classmethod
  123. def poll(cls, ntree):
  124. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  125. @classmethod
  126. def set_mantis_class(self):
  127. from importlib import import_module
  128. # do not catch errors, they should cause a failure.
  129. try:
  130. module = import_module(self.mantis_node_library, package=mantis_root)
  131. self.mantis_class=getattr(module, self.mantis_node_class_name)
  132. except Exception as e:
  133. print(self)
  134. raise e
  135. def insert_link(self, link):
  136. if (bpy.app.version < (4, 4, 0)):
  137. return # this causes a crasah due to a bug.
  138. context = bpy.context
  139. if context.space_data:
  140. node_tree = context.space_data.path[0].node_tree
  141. if node_tree.do_live_update:
  142. node_tree.update_tree(context)
  143. if (link.to_socket.is_linked == False):
  144. node_tree.num_links+=1
  145. elif (link.to_socket.is_multi_input):
  146. node_tree.num_links+=1
  147. class SchemaUINode(MantisUINode):
  148. mantis_node_library='.schema_containers'
  149. @classmethod
  150. def poll(cls, ntree):
  151. return (ntree.bl_idname in ['SchemaTree'])
  152. class LinkNode(MantisUINode):
  153. mantis_node_library='.link_containers'
  154. @classmethod
  155. def poll(cls, ntree):
  156. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  157. class xFormNode(MantisUINode):
  158. mantis_node_library='.xForm_containers'
  159. @classmethod
  160. def poll(cls, ntree):
  161. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  162. class DeformerNode(MantisUINode):
  163. mantis_node_library='.deformer_containers'
  164. @classmethod
  165. def poll(cls, ntree):
  166. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  167. def poll_node_tree(self, object):
  168. if isinstance(object, MantisTree):
  169. return True
  170. return False
  171. # TODO: try to check identifiers instead of name.
  172. def node_group_update(node, force = False):
  173. if not force:
  174. if (node.id_data.do_live_update == False) or (node.id_data.is_executing == True):
  175. return
  176. # note: if (node.id_data.is_exporting == True) I need to be able to update so I can make links.
  177. toggle_update = node.id_data.do_live_update
  178. node.id_data.do_live_update = False
  179. identifiers_in={socket.identifier:socket for socket in node.inputs}
  180. identifiers_out={socket.identifier:socket for socket in node.outputs}
  181. if node.node_tree is None:
  182. node.inputs.clear(); node.outputs.clear()
  183. node.id_data.do_live_update = toggle_update
  184. return
  185. found_in, found_out = [], []
  186. update_input, update_output = False, False
  187. for item in node.node_tree.interface.items_tree:
  188. if item.item_type != "SOCKET": continue
  189. if item.in_out == 'OUTPUT':
  190. if s:= identifiers_out.get(item.identifier): # if the requested output doesn't exist, update
  191. found_out.append(item.identifier)
  192. if update_output: continue
  193. if s.bl_idname != item.socket_type: update_output = True; continue
  194. else: update_output = True; continue
  195. else:
  196. if s:= identifiers_in.get(item.identifier): # if the requested input doesn't exist, update
  197. found_in.append(item.identifier)
  198. if update_input: continue # done here
  199. if s.bl_idname != item.socket_type: update_input = True; continue
  200. else: update_input = True; continue
  201. # Schema has an extra input for Length and for Extend.
  202. if node.bl_idname == 'MantisSchemaGroup':
  203. found_in.extend(['Schema Length', ''])
  204. # if we have too many elements, just get rid of the ones we don't need
  205. if len(node.inputs) > len(found_in):#
  206. for inp in node.inputs:
  207. if inp.identifier in found_in: continue
  208. node.inputs.remove(inp)
  209. if len(node.outputs) > len(found_out):
  210. for out in node.outputs:
  211. if out.identifier in found_out: continue
  212. node.outputs.remove(out)
  213. #
  214. if len(node.inputs) > 0 and (inp := node.inputs[-1]).bl_idname == 'WildcardSocket' and inp.is_linked:
  215. update_input = True
  216. if len(node.outputs) > 0 and (out := node.outputs[-1]).bl_idname == 'WildcardSocket' and out.is_linked:
  217. update_output = True
  218. #
  219. if not (update_input or update_output):
  220. node.id_data.do_live_update = toggle_update
  221. return
  222. if update_input or update_output:
  223. socket_map_in, socket_map_out = get_socket_maps(node)
  224. if update_input :
  225. if node.bl_idname == 'MantisSchemaGroup':
  226. schema_length=0
  227. if sl := node.inputs.get("Schema Length"):
  228. schema_length = sl.default_value
  229. # sometimes this isn't available yet # TODO not happy about this solution
  230. node.inputs.clear()
  231. if node.bl_idname == 'MantisSchemaGroup':
  232. node.inputs.new("IntSocket", "Schema Length", identifier='Schema Length')
  233. node.inputs['Schema Length'].default_value = schema_length
  234. if update_output: node.outputs.clear()
  235. for item in node.node_tree.interface.items_tree:
  236. if item.item_type != "SOCKET": continue
  237. if (item.in_out == 'INPUT' and update_input):
  238. relink_socket_map(node, node.inputs, socket_map_in, item)
  239. if (item.in_out == 'OUTPUT' and update_output):
  240. relink_socket_map(node, node.outputs, socket_map_out, item)
  241. # at this point there is no wildcard socket
  242. if '__extend__' in socket_map_in.keys():
  243. do_relink(node, None, socket_map_in, in_out='INPUT', parent_name='Constant' )
  244. node.id_data.do_live_update = toggle_update
  245. def node_tree_prop_update(self, context):
  246. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  247. return # so we check if an update is currently running.
  248. self.is_updating = True
  249. node_group_update(self)
  250. self.is_updating = False
  251. if self.bl_idname in ['MantisSchemaGroup'] and self.node_tree is not None:
  252. if len(self.inputs) == 0:
  253. self.inputs.new("IntSocket", "Schema Length", identifier='Schema Length')
  254. if self.inputs[-1].bl_idname != "WildcardSocket":
  255. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  256. from bpy.types import NodeCustomGroup
  257. class MantisNodeGroup(Node, MantisUINode):
  258. bl_idname = "MantisNodeGroup"
  259. bl_label = "Node Group"
  260. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree, update=node_tree_prop_update,)
  261. is_updating:BoolProperty(default=False)
  262. def update(self):
  263. live_update = self.id_data.do_live_update
  264. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  265. return # so we check if an update is currently running.
  266. try:
  267. self.is_updating = True
  268. node_group_update(self)
  269. self.is_updating = False
  270. finally: # we need to reset this regardless of whether or not the operation succeeds!
  271. self.is_updating = False
  272. self.id_data.do_live_update = live_update # ensure this remains the same
  273. def draw_buttons(self, context, layout):
  274. row = layout.row(align=True)
  275. row.prop(self, "node_tree", text="")
  276. row.operator("mantis.edit_group", text="", icon='NODETREE', emboss=True)
  277. class GraphError(Exception):
  278. pass
  279. def get_signature_from_edited_tree(node, context):
  280. sig_path=[None,]
  281. for item in context.space_data.path[:-1]:
  282. sig_path.append(item.node_tree.nodes.active.name)
  283. return tuple(sig_path+[node.name])
  284. def poll_node_tree_schema(self, object):
  285. if isinstance(object, SchemaTree):
  286. return True
  287. return False
  288. # TODO tiny UI problem - inserting new links into the tree will not place them in the right place.
  289. class SchemaGroup(Node, MantisUINode):
  290. bl_idname = "MantisSchemaGroup"
  291. bl_label = "Node Schema"
  292. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree_schema, update=node_tree_prop_update,)
  293. is_updating:BoolProperty(default=False)
  294. def draw_buttons(self, context, layout):
  295. row = layout.row(align=True)
  296. row.prop(self, "node_tree", text="")
  297. row.operator("mantis.edit_group", text="", icon='NODETREE', emboss=True)
  298. def update(self):
  299. live_update = self.id_data.do_live_update
  300. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  301. return # so we check if an update is currently running.
  302. self.is_updating = True
  303. try:
  304. node_group_update(self)
  305. # reset things if necessary:
  306. if self.node_tree:
  307. if len(self.inputs) == 0:
  308. self.inputs.new("IntSocket", "Schema Length", identifier='Schema Length')
  309. if self.inputs[-1].bl_idname != "WildcardSocket":
  310. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  311. finally: # we need to reset this regardless of whether or not the operation succeeds!
  312. self.is_updating = False
  313. self.id_data.do_live_update = live_update # ensure this remains the same
  314. NODES_REMOVED=["xFormRootNode"]
  315. # Node bl_idname, # Socket Name
  316. SOCKETS_REMOVED=[("UtilityDriverVariable", "Transform Channel"),
  317. ("xFormRootNode","World Out"),
  318. ("UtilitySwitch","xForm"),
  319. ("LinkDrivenParameter", "Enable")]
  320. # Node Class #Prior bl_idname # prior name # new bl_idname # new name, # Multi
  321. SOCKETS_RENAMED=[ ("LinkDrivenParameter", "DriverSocket", "Driver", "FloatSocket", "Value", False)]
  322. # NODE CLASS NAME IN_OUT SOCKET TYPE SOCKET NAME INDEX MULTI DEFAULT
  323. SOCKETS_ADDED=[("DeformerMorphTargetDeform", 'INPUT', 'BooleanSocket', "Use Shape Key", 1, False, False),
  324. ("DeformerMorphTargetDeform", 'INPUT', 'BooleanSocket', "Use Offset", 2, False, True),
  325. ("UtilityFCurve", 'INPUT', "eFCrvExtrapolationMode", "Extrapolation Mode", 0, False, 'CONSTANT')]
  326. # replace names with bl_idnames for reading the tree and solving schemas.
  327. replace_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",
  328. "SchemaArrayInput", "SchemaConstInput", "SchemaConstOutput", "SchemaIndex",
  329. "SchemaOutgoingConnection", "SchemaConstantOutput", "SchemaArrayOutput",
  330. "SchemaArrayInputGet",]
  331. # anything that gets properties added in the graph... this is a clumsy approach but I need to watch for this
  332. # in schema generation and this is the easiest way to do it for now.
  333. custom_props_types = ["LinkArmature", "UtilityKeyframe", "UtilityFCurve", "UtilityDriver", "xFormBone"]
  334. # filters for determining if a link is a hierarchy link or a non-hierarchy (cyclic) link.
  335. from_name_filter = ["Driver",]
  336. to_name_filter = [
  337. "Custom Object xForm Override",
  338. "Custom Object",
  339. "Deform Bones",
  340. ]
  341. class MantisNode:
  342. """
  343. This class contains the basic interface for a Mantis Node.
  344. A MantisNode is used internally by Mantis to represent the final evaluated node graph.
  345. It gets generated with data from a MantisUINode when the graph is read.
  346. """
  347. def __init__(self, signature, base_tree):
  348. self.base_tree=base_tree
  349. self.signature = signature
  350. self.inputs = MantisNodeSocketCollection(node=self, is_input=True)
  351. self.outputs = MantisNodeSocketCollection(node=self, is_input=False)
  352. self.parameters = {}
  353. self.drivers = {}
  354. self.node_type='UNINITIALIZED'
  355. self.hierarchy_connections, self.connections = [], []
  356. self.hierarchy_dependencies, self.dependencies = [], []
  357. self.prepared = False
  358. self.executed = False
  359. def init_parameters(self, additional_parameters = {}):
  360. for socket in self.inputs:
  361. self.parameters[socket.name] = None
  362. for socket in self.outputs:
  363. self.parameters[socket.name] = None
  364. for key, value in additional_parameters.items():
  365. self.parameters[key]=value
  366. def set_traverse(self, traversal_pairs = [(str, str)]):
  367. for (a, b) in traversal_pairs:
  368. self.inputs[a].set_traverse_target(self.outputs[b])
  369. self.outputs[b].set_traverse_target(self.inputs[a])
  370. def flush_links(self):
  371. for inp in self.inputs.values():
  372. inp.flush_links()
  373. for out in self.outputs.values():
  374. out.flush_links()
  375. def evaluate_input(self, input_name, index=0):
  376. from .node_container_common import trace_single_line
  377. if not (self.inputs.get(input_name)): # get the named parameter if there is no input
  378. return self.parameters.get(input_name) # this will return None if the parameter does not exist.
  379. # this trace() should give a key error if there is a problem
  380. # it is NOT handled here because it should NOT happen - so I want the error message.
  381. trace = trace_single_line(self, input_name, index)
  382. prop = trace[0][-1].parameters[trace[1].name] #trace[0] = the list of traced nodes; read its parameters
  383. return prop
  384. def fill_parameters(self, ui_node=None):
  385. from .utilities import get_node_prototype
  386. from .node_container_common import get_socket_value
  387. if not ui_node:
  388. if ( (self.signature[0] in ["MANTIS_AUTOGENERATED", "SCHEMA_AUTOGENERATED" ]) or
  389. (self.signature[-1] in ["NodeGroupOutput", "NodeGroupInput"]) ): # I think this is harmless
  390. return None
  391. else:
  392. ui_node = get_node_prototype(self.signature, self.base_tree)
  393. if not ui_node:
  394. raise RuntimeError(wrapRed("No node prototype found for... %s" % ( [self.base_tree] + list(self.signature[1:]) ) ) )
  395. for key in self.parameters.keys():
  396. node_socket = ui_node.inputs.get(key)
  397. if self.parameters[key] is not None: # the parameters are usually initialized as None.
  398. continue # will be filled by the node itself
  399. if not node_socket: #maybe the node socket has no name
  400. if ( ( len(ui_node.inputs) == 0) and ( len(ui_node.outputs) == 1) ):
  401. # this is a simple input node.
  402. node_socket = ui_node.outputs[0]
  403. elif key == 'Name': # for Links we just use the Node Label, or if there is no label, the name.
  404. self.parameters[key] = ui_node.label if ui_node.label else ui_node.name
  405. continue
  406. else:
  407. pass
  408. if node_socket:
  409. if node_socket.bl_idname in ['RelationshipSocket', 'xFormSocket']:
  410. continue
  411. elif node_socket.is_linked and (not node_socket.is_output):
  412. pass # we will get the value from the link, because this is a linked input port.
  413. # very importantly, we do not pass linked outputs- fill these because they are probably Input nodes.
  414. elif hasattr(node_socket, "default_value"):
  415. if (value := get_socket_value(node_socket)) is not None:
  416. self.parameters[key] = value
  417. # TODO: try and remove the input if it is not needed (for performance speed)
  418. else:
  419. raise RuntimeError(wrapRed("No value found for " + self.__repr__() + " when filling out node parameters for " + ui_node.name + "::"+node_socket.name))
  420. else:
  421. pass
  422. def call_on_all_ancestors(self, *args, **kwargs):
  423. """Resolve the dependencies of this node with the named method and its arguments.
  424. First, dependencies are discovered by walking backwards through the tree. Once the root
  425. nodes are discovered, the method is called by each node in dependency order.
  426. The first argument MUST be the name of the method as a string.
  427. """
  428. prGreen(self)
  429. if args[0] == 'call_on_all_ancestors': raise RuntimeError("Very funny!")
  430. from .utilities import get_all_dependencies
  431. from collections import deque
  432. # get all dependencies by walking backward through the tree.
  433. all_dependencies = get_all_dependencies(self)
  434. # get just the roots
  435. can_solve = deque(filter(lambda a : len(a.hierarchy_connections) == 0, all_dependencies))
  436. solved = set()
  437. while can_solve:
  438. node = can_solve.pop()
  439. print(node)
  440. method = getattr(node, args[0])
  441. method(*args[0:], **kwargs)
  442. solved.add(node)
  443. can_solve.extendleft(filter(lambda a : a in all_dependencies, node.hierarchy_connections))
  444. # prPurple(can_solve)
  445. if self in solved:
  446. break
  447. # else:
  448. # for dep in all_dependencies:
  449. # if dep not in solved:
  450. # prOrange(dep)
  451. return
  452. def bPrepare(self, bContext=None):
  453. return
  454. def bExecute(self, bContext=None):
  455. return
  456. def bFinalize(self, bContext=None):
  457. return
  458. def __repr__(self):
  459. return self.signature.__repr__()
  460. # do I need this and the link class above?
  461. class DummyLink:
  462. #gonna use this for faking links to keep the interface consistent
  463. def __init__(self, from_socket, to_socket, nc_from=None, nc_to=None, original_from=None, multi_input_sort_id=0):
  464. self.from_socket = from_socket
  465. self.to_socket = to_socket
  466. self.nc_from = nc_from
  467. self.nc_to = nc_to
  468. self.multi_input_sort_id = multi_input_sort_id
  469. # self.from_node = from_socket.node
  470. # self.to_node = to_socket.node
  471. if (original_from):
  472. self.original_from = original_from
  473. else:
  474. self.original_from = self.from_socket
  475. def __repr__(self):
  476. return(self.nc_from.__repr__()+":"+self.from_socket.name + " -> " + self.nc_to.__repr__()+":"+self.to_socket.name)
  477. class NodeLink:
  478. from_node = None
  479. from_socket = None
  480. to_node = None
  481. to_socket = None
  482. def __init__(self, from_node, from_socket, to_node, to_socket, multi_input_sort_id=0):
  483. if from_node.signature == to_node.signature:
  484. raise RuntimeError("Cannot connect a node to itself.")
  485. self.from_node = from_node
  486. self.from_socket = from_socket
  487. self.to_node = to_node
  488. self.to_socket = to_socket
  489. self.from_node.outputs[self.from_socket].links.append(self)
  490. # it is the responsibility of the node that uses these links to sort them correctly based on the sort_id
  491. self.multi_input_sort_id = multi_input_sort_id
  492. self.to_node.inputs[self.to_socket].links.append(self)
  493. from .node_container_common import detect_hierarchy_link
  494. self.is_hierarchy = detect_hierarchy_link(from_node, from_socket, to_node, to_socket,)
  495. self.is_alive = True
  496. def __repr__(self):
  497. return self.from_node.outputs[self.from_socket].__repr__() + " --> " + self.to_node.inputs[self.to_socket].__repr__()
  498. # link_string = # if I need to colorize output for debugging.
  499. # if self.is_hierarchy:
  500. # return wrapOrange(link_string)
  501. # else:
  502. # return wrapWhite(link_string)
  503. def die(self):
  504. self.is_alive = False
  505. self.to_node.inputs[self.to_socket].flush_links()
  506. self.from_node.outputs[self.from_socket].flush_links()
  507. def insert_node(self, middle_node, middle_node_in, middle_node_out, re_init_hierarchy = True):
  508. to_node = self.to_node
  509. to_socket = self.to_socket
  510. self.to_node = middle_node
  511. self.to_socket = middle_node_in
  512. middle_node.outputs[middle_node_out].connect(to_node, to_socket)
  513. if re_init_hierarchy:
  514. from .utilities import init_connections, init_dependencies
  515. init_connections(self.from_node)
  516. init_connections(middle_node)
  517. init_dependencies(middle_node)
  518. init_dependencies(to_node)
  519. class NodeSocket:
  520. # @property # this is a read-only property.
  521. # def is_linked(self):
  522. # return bool(self.links)
  523. def __init__(self, is_input = False,
  524. node = None, name = None,
  525. traverse_target = None):
  526. self.can_traverse = False # to/from the other side of the parent node
  527. self.traverse_target = None
  528. self.node = node
  529. self.name = name
  530. self.is_input = is_input
  531. self.links = []
  532. self.is_linked = False
  533. if (traverse_target):
  534. self.can_traverse = True
  535. def connect(self, node, socket, sort_id=0):
  536. if (self.is_input):
  537. to_node = self.node; from_node = node
  538. to_socket = self.name; from_socket = socket
  539. else:
  540. from_node = self.node; to_node = node
  541. from_socket = self.name; to_socket = socket
  542. from_node.outputs[from_socket].is_linked = True
  543. to_node.inputs[to_socket].is_linked = True
  544. for l in from_node.outputs[from_socket].links:
  545. if l.to_node==to_node and l.to_socket==to_socket:
  546. return None
  547. new_link = NodeLink(
  548. from_node,
  549. from_socket,
  550. to_node,
  551. to_socket,
  552. sort_id)
  553. return new_link
  554. def set_traverse_target(self, traverse_target):
  555. self.traverse_target = traverse_target
  556. self.can_traverse = True
  557. def flush_links(self):
  558. """ Removes dead links from this socket."""
  559. self.links = [l for l in self.links if l.is_alive]
  560. self.is_linked = bool(self.links)
  561. @property
  562. def is_connected(self):
  563. return len(self.links)>0
  564. def __repr__(self):
  565. return self.node.__repr__() + "::" + self.name
  566. class MantisNodeSocketCollection(dict):
  567. def __init__(self, node, is_input=False):
  568. self.is_input = is_input
  569. self.node = node
  570. def init_sockets(self, sockets):
  571. for socket in sockets:
  572. if not isinstance(socket, str): raise RuntimeError("NodeSocketCollection keys must be str.")
  573. self[socket] = NodeSocket(is_input=self.is_input, name=socket, node=self.node)
  574. def __delitem__(self, key):
  575. """Deletes a node socket by name, and all its links."""
  576. socket = self[key]
  577. for l in socket.links:
  578. l.die()
  579. super().__delitem__(key)
  580. def __iter__(self):
  581. """Makes the class iterable"""
  582. return iter(self.values())