base_definitions.py 51 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175
  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. FLOAT_EPSILON=0.0001 # used to check against floating point inaccuracy
  13. links_sort_key= lambda a : (-a.multi_input_sort_id, -a.sub_sort_id)
  14. def TellClasses():
  15. #Why use a function to do this? Because I don't need every class to register.
  16. return [ MantisTree,
  17. SchemaTree,
  18. MantisNodeGroup,
  19. SchemaGroup,
  20. ]
  21. def error_popup_draw(self, context):
  22. self.layout.label(text="Error executing tree. See Console.")
  23. mantis_root = ".".join(__name__.split('.')[:-1]) # absolute HACK
  24. from os import environ
  25. # https://docs.blender.org/api/master/bpy.types.NodeTree.html#bpy.types.NodeTree.valid_socket_type
  26. # thank you, Sverchok
  27. def valid_interface_types(cls : NodeTree, socket_idname : str):
  28. from .socket_definitions import tell_valid_bl_idnames, TellClasses
  29. #TODO: do the versioning code to handle this so it can be in all versions
  30. if bpy.app.version <= (4,4,0): # should work in 4.4.1
  31. return socket_idname in [cls.bl_idname for cls in TellClasses()]
  32. elif bpy.app.version == (4,5,0): # workaround for a BUG
  33. return ['NodeSocketGeometry']
  34. else: # once versioning is finished this will be unnecesary.
  35. return socket_idname in tell_valid_bl_idnames()
  36. def fix_reroute_colors(tree):
  37. context = bpy.context
  38. if any((tree.is_executing, tree.is_exporting, tree.do_live_update==False, context.space_data is None) ):
  39. return
  40. from collections import deque
  41. from .utilities import socket_seek
  42. from .socket_definitions import MantisSocket
  43. reroutes_without_color = deque()
  44. for n in tree.nodes:
  45. if n.bl_idname=='NodeReroute' and n.inputs[0].bl_idname == "NodeSocketColor":
  46. reroutes_without_color.append(n)
  47. try:
  48. while reroutes_without_color:
  49. rr = reroutes_without_color.pop()
  50. if rr.inputs[0].is_linked:
  51. link = rr.inputs[0].links[0]
  52. socket = socket_seek(link, tree.links)
  53. if isinstance(socket, MantisSocket):
  54. rr.socket_idname = socket.bl_idname
  55. except Exception as e:
  56. print(wrapOrange("WARN: Updating reroute color failed with exception: ")+wrapWhite(f"{e.__class__.__name__}"))
  57. #functions to identify the state of the system using hashes
  58. # this function runs a lot so it should be optimized as well as possible.
  59. def hash_tree(tree):
  60. trees=set(); links=[]; hash_data=""
  61. for node in tree.nodes:
  62. hash_data+=str(node.name)
  63. if hasattr(node, 'node_tree') and node.node_tree:
  64. trees.add(node.node_tree)
  65. for other_tree in trees:
  66. hash_data+=str(hash_tree(other_tree))
  67. for link in tree.links:
  68. links.append( link.from_node.name+link.from_socket.name+
  69. link.to_node.name+link.to_socket.name+
  70. str(link.multi_input_sort_id) )
  71. links.sort(); hash_data+=''.join(links)
  72. return hash(hash_data)
  73. MANTIS_VERSION_MAJOR=0
  74. MANTIS_VERSION_MINOR=13
  75. MANTIS_VERSION_SUB=0
  76. class MantisTree(NodeTree):
  77. '''A custom node tree type that will show up in the editor type list'''
  78. bl_idname = 'MantisTree'
  79. bl_label = "Rigging Nodes"
  80. bl_icon = 'OUTLINER_OB_ARMATURE'
  81. tree_valid:BoolProperty(default=False)
  82. hash:StringProperty(default='')
  83. do_live_update:BoolProperty(default=True) # use this to disable updates for e.g. scripts
  84. num_links:IntProperty(default=-1)
  85. # operator settings for re-exporting the tree.
  86. filepath:StringProperty(default="", subtype='FILE_PATH')
  87. export_all_subtrees_together:BoolProperty(default=True)
  88. #
  89. is_executing:BoolProperty(default=False)
  90. is_exporting:BoolProperty(default=False)
  91. execution_id:StringProperty(default='')
  92. # prev_execution_id:StringProperty(default='')
  93. mantis_version:IntVectorProperty(default=[
  94. MANTIS_VERSION_MAJOR, MANTIS_VERSION_MINOR, MANTIS_VERSION_SUB])
  95. # this prevents the node group from executing on the next depsgraph update
  96. # because I don't always have control over when the dg upadte happens.
  97. prevent_next_exec:BoolProperty(default=False)
  98. #added to work around a bug in 4.5.0 LTS
  99. interface_helper : StringProperty(default='')
  100. parsed_tree={}
  101. if (bpy.app.version < (4, 4, 0) or bpy.app.version >= (4,5,0)): # in 4.4 this leads to a crash
  102. @classmethod
  103. def valid_socket_type(cls : NodeTree, socket_idname: str):
  104. return valid_interface_types(cls, socket_idname)
  105. def update(self): # set the reroute colors
  106. if (bpy.app.version >= (4,4,0)):
  107. fix_reroute_colors(self)
  108. def update_tree(self, context = None, force=False, error_popups=False):
  109. if self.is_exporting:
  110. return
  111. my_hash = str( hash_tree(self) )
  112. if (my_hash != self.hash) or force:
  113. self.hash = my_hash
  114. self.is_executing = True
  115. from . import readtree
  116. prGreen("Validating Tree: %s" % self.name)
  117. try:
  118. import bpy # I am importing here so that the context passed in
  119. # is used for display update... but I always want to do this
  120. scene = bpy.context.scene
  121. scene.render.use_lock_interface = True
  122. self.parsed_tree = readtree.parse_tree(self, error_popups)
  123. if context:
  124. self.display_update(context)
  125. self.tree_valid = True
  126. except Exception as e:
  127. prRed("Failed to update node tree due to error.")
  128. self.tree_valid = False
  129. self.hash='' # unset the hash to mark the tree as un-parsed.
  130. raise e
  131. finally:
  132. scene.render.use_lock_interface = False
  133. self.is_executing = False
  134. def display_update(self, context):
  135. if self.is_exporting:
  136. return
  137. self.is_executing = True
  138. current_tree = bpy.context.space_data.path[-1].node_tree
  139. for node in current_tree.nodes:
  140. if hasattr(node, "display_update"):
  141. try:
  142. node.display_update(self.parsed_tree, context)
  143. except Exception as e:
  144. print("Node \"%s\" failed to update display with error: %s" %(wrapGreen(node.name), wrapRed(e)))
  145. self.is_executing = False
  146. # TODO: deal with invalid links properly.
  147. # - Non-hierarchy links should be ignored in the circle-check and so the links should be marked valid in such a circle
  148. # - hierarchy-links should be marked invalid and prevent the tree from executing.
  149. def execute_tree(self,context, error_popups = False):
  150. self.prevent_next_exec = False
  151. if not self.hash:
  152. return
  153. if self.is_exporting or self.is_executing:
  154. return
  155. prGreen("Executing Tree: %s" % self.name)
  156. self.is_executing = True
  157. from . import readtree
  158. try:
  159. context.scene.render.use_lock_interface = True
  160. readtree.execute_tree(self.parsed_tree, self, context, error_popups)
  161. except RecursionError as e:
  162. prRed("Recursion error while parsing tree.")
  163. finally:
  164. context.scene.render.use_lock_interface = False
  165. self.is_executing = False
  166. class SchemaTree(NodeTree):
  167. '''A node tree representing a schema to generate a Mantis tree'''
  168. bl_idname = 'SchemaTree'
  169. bl_label = "Rigging Nodes Schema"
  170. bl_icon = 'RIGID_BODY_CONSTRAINT'
  171. # these are only needed for consistent interface, but should not be used
  172. do_live_update:BoolProperty(default=True) # default to true so that updates work
  173. is_executing:BoolProperty(default=False)
  174. is_exporting:BoolProperty(default=False)
  175. mantis_version:IntVectorProperty(default=[
  176. MANTIS_VERSION_MAJOR, MANTIS_VERSION_MINOR, MANTIS_VERSION_SUB])
  177. # see the note in MantisTree
  178. interface_helper : StringProperty(default='')
  179. if (bpy.app.version < (4, 4, 0) or bpy.app.version >= (4,5,0)): # in 4.4 this leads to a crash
  180. @classmethod
  181. def valid_socket_type(cls : NodeTree, socket_idname: str):
  182. return valid_interface_types(cls, socket_idname)
  183. def update(self): # set the reroute colors
  184. if (bpy.app.version >= (4,4,0)):
  185. fix_reroute_colors(self)
  186. from dataclasses import dataclass, field
  187. from typing import Any
  188. @dataclass
  189. class MantisSocketTemplate():
  190. name : str = field(default="")
  191. bl_idname : str = field(default="")
  192. traverse_target : str = field(default="")
  193. identifier : str = field(default="")
  194. display_shape : str = field(default="") # for arrays
  195. category : str = field(default="") # for use in display update
  196. blender_property : str | tuple[str] = field(default="") # for props_sockets -> evaluate sockets
  197. is_input : bool = field(default=False)
  198. hide : bool = field(default=False)
  199. use_multi_input : bool = field(default=False)
  200. default_value : Any = field(default=None)
  201. #TODO: do a better job explaining how MantisNode and MantisUINode relate.
  202. class MantisUINode:
  203. """
  204. This class contains the common user-interface features of Mantis nodes.
  205. MantisUINode objects will spawn one or more MantisNode objects when the graph is evaluated.
  206. The MantisNode objects will pull the data from the UI node and use it to generate the graph.
  207. """
  208. mantis_node_library=''
  209. mantis_node_class_name=''
  210. mantis_class=None
  211. @classmethod
  212. def poll(cls, ntree):
  213. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  214. @classmethod
  215. def set_mantis_class(self):
  216. from importlib import import_module
  217. # do not catch errors, they should cause a failure.
  218. try:
  219. module = import_module(self.mantis_node_library, package=mantis_root)
  220. self.mantis_class=getattr(module, self.mantis_node_class_name)
  221. except Exception as e:
  222. print(self)
  223. raise e
  224. def insert_link(self, link):
  225. if (bpy.app.version >= (4, 4, 0)):
  226. return # this causes a crash due to a bug.
  227. context = bpy.context
  228. if context.space_data:
  229. node_tree = context.space_data.path[0].node_tree
  230. if node_tree.do_live_update:
  231. node_tree.update_tree(context)
  232. if (link.to_socket.is_linked == False):
  233. node_tree.num_links+=1
  234. elif (link.to_socket.is_multi_input):
  235. node_tree.num_links+=1
  236. def init_sockets(self, socket_templates : tuple[MantisSocketTemplate]):
  237. for template in socket_templates:
  238. collection = self.outputs
  239. if template.is_input:
  240. collection = self.inputs
  241. identifier = template.name
  242. if template.identifier:
  243. identifier = template.identifier
  244. use_multi_input = template.use_multi_input if template.is_input else False
  245. socket = collection.new(
  246. template.bl_idname,
  247. template.name,
  248. identifier=identifier,
  249. use_multi_input=use_multi_input
  250. )
  251. socket.hide= template.hide
  252. if template.category:
  253. # a custom property for the UI functions to use.
  254. socket['category'] = template.category
  255. if template.default_value is not None:
  256. socket.default_value = template.default_value
  257. # this can throw a TypeError - it is the caller's
  258. # responsibility to send the right type.
  259. if template.use_multi_input: # this is an array
  260. socket.display_shape = 'SQUARE_DOT'
  261. class SchemaUINode(MantisUINode):
  262. mantis_node_library='.schema_nodes'
  263. is_updating:BoolProperty(default=False)
  264. @classmethod
  265. def poll(cls, ntree):
  266. return (ntree.bl_idname in ['SchemaTree'])
  267. class LinkNode(MantisUINode):
  268. mantis_node_library='.link_nodes'
  269. @classmethod
  270. def poll(cls, ntree):
  271. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  272. class xFormNode(MantisUINode):
  273. mantis_node_library='.xForm_nodes'
  274. @classmethod
  275. def poll(cls, ntree):
  276. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  277. class DeformerNode(MantisUINode):
  278. mantis_node_library='.deformer_nodes'
  279. @classmethod
  280. def poll(cls, ntree):
  281. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  282. def poll_node_tree(self, object):
  283. forbid = []
  284. context = bpy.context
  285. if context.space_data:
  286. if context.space_data.path:
  287. for path_item in context.space_data.path:
  288. forbid.append(path_item.node_tree.name)
  289. if isinstance(object, MantisTree) and object.name not in forbid:
  290. return True
  291. return False
  292. # TODO: try and remove the extra loop used here... but it is OK for now
  293. def should_remove_socket(node, socket):
  294. # a function to check if the socket is in the interface
  295. id_found = False
  296. for item in node.node_tree.interface.items_tree:
  297. if item.item_type != "SOCKET": continue
  298. if item.identifier == socket.identifier:
  299. id_found = True; break
  300. return not id_found
  301. # TODO: try to check identifiers instead of name.
  302. def node_group_update(node, force = False):
  303. if not node.is_updating:
  304. raise RuntimeError("Cannot update node while it is not marked as updating.")
  305. if not force:
  306. if (node.id_data.do_live_update == False) or \
  307. (node.id_data.is_executing == True) or \
  308. (node.id_data.is_exporting == True):
  309. return
  310. # note: if (node.id_data.is_exporting == True) I need to be able to update so I can make links.
  311. if node.node_tree is None:
  312. node.inputs.clear(); node.outputs.clear()
  313. return
  314. toggle_update = node.id_data.do_live_update
  315. node.id_data.do_live_update = False
  316. identifiers_in={socket.identifier:socket for socket in node.inputs}
  317. identifiers_out={socket.identifier:socket for socket in node.outputs}
  318. interface_names_in, interface_names_out = {}, {}
  319. for interface_item in node.node_tree.interface.items_tree:
  320. if interface_item.item_type != 'SOCKET': continue
  321. if interface_item.in_out == 'INPUT':
  322. interface_names_in[interface_item.identifier] = interface_item.name
  323. else:
  324. interface_names_out[interface_item.identifier] = interface_item.name
  325. indices_in,indices_out={},{} # check by INDEX to see if the socket's name/type match.
  326. for collection, map in [(node.inputs, indices_in), (node.outputs, indices_out)]:
  327. for i, socket in enumerate(collection):
  328. map[socket.identifier]=i
  329. if node.node_tree is None:
  330. node.inputs.clear(); node.outputs.clear()
  331. node.id_data.do_live_update = toggle_update
  332. return
  333. found_in, found_out = [], []
  334. update_input, update_output = False, False
  335. for item in node.node_tree.interface.items_tree:
  336. if item.item_type != "SOCKET": continue
  337. if item.in_out == 'OUTPUT':
  338. if s:= identifiers_out.get(item.identifier): # if the requested output doesn't exist, update
  339. if interface_names_out.get(item.identifier) != s.name: update_output = True; continue
  340. found_out.append(item.identifier)
  341. if (indices_out[s.identifier]!=item.index): update_output=True; continue
  342. if update_output: continue
  343. if s.bl_idname != item.bl_socket_idname: update_output = True; continue
  344. else: update_output = True; continue
  345. else:
  346. if s:= identifiers_in.get(item.identifier): # if the requested input doesn't exist, update
  347. if interface_names_in.get(item.identifier) != s.name: update_input = True; continue
  348. found_in.append(item.identifier)
  349. if (indices_in[s.identifier]!=item.index): update_input=True; continue
  350. if update_input: continue # done here
  351. if s.bl_idname != item.bl_socket_idname: update_input = True; continue
  352. else: update_input = True; continue
  353. # Schema has an extra input for Length and for Extend.
  354. if node.bl_idname == 'MantisSchemaGroup':
  355. found_in.extend(['Schema Length', ''])
  356. # get the socket maps before modifying stuff
  357. if update_input or update_output:
  358. socket_maps = get_socket_maps(node,)
  359. if socket_maps:
  360. socket_map_in, socket_map_out = socket_maps
  361. if node.bl_idname == "MantisSchemaGroup" and \
  362. len(node.inputs)+len(node.outputs)<=2 and\
  363. len(node.node_tree.interface.items_tree) > 0:
  364. socket_map_in, socket_map_out = None, None
  365. # We have to initialize the node because it only has its base inputs.
  366. elif socket_maps is None:
  367. node.id_data.do_live_update = toggle_update
  368. # if we have too many elements, just get rid of the ones we don't need
  369. if len(node.inputs) > len(found_in):#
  370. for inp in node.inputs:
  371. if inp.identifier in found_in: continue
  372. node.inputs.remove(inp)
  373. if len(node.outputs) > len(found_out):
  374. for out in node.outputs:
  375. if out.identifier in found_out: continue
  376. node.outputs.remove(out)
  377. #
  378. if len(node.inputs) > 0 and (inp := node.inputs[-1]).bl_idname == 'WildcardSocket' and inp.is_linked:
  379. update_input = True
  380. #
  381. if not (update_input or update_output):
  382. node.id_data.do_live_update = toggle_update
  383. return
  384. if update_input or update_output and (socket_maps is not None):
  385. if update_input :
  386. if node.bl_idname == 'MantisSchemaGroup':
  387. schema_length=0
  388. if sl := node.inputs.get("Schema Length"):
  389. schema_length = sl.default_value
  390. # sometimes this isn't available yet # TODO not happy about this solution
  391. remove_me=[]
  392. # remove all found map items but the Schema Length input (reuse it)
  393. for i, socket in enumerate(node.inputs):
  394. if socket.identifier == "Schema Length" and i == 0:
  395. continue
  396. elif (socket_map_in is None) or socket.identifier in socket_map_in.keys():
  397. remove_me.append(socket)
  398. elif should_remove_socket(node, socket):
  399. remove_me.append(socket)
  400. while remove_me:
  401. node.inputs.remove(remove_me.pop())
  402. if update_output:
  403. remove_me=[]
  404. for socket in node.outputs:
  405. if (socket_map_out is None) or socket.identifier in socket_map_out.keys():
  406. remove_me.append(socket)
  407. elif should_remove_socket(node, socket):
  408. remove_me.append(socket)
  409. while remove_me:
  410. node.inputs.remove(remove_me.pop())
  411. from .utilities import relink_socket_map_add_socket
  412. reorder_me_input = []; input_index = 0
  413. reorder_me_output = []; output_index = 0
  414. def update_group_sockets(interface_item, is_input):
  415. socket_map = socket_map_in if is_input else socket_map_out
  416. socket_collection = node.inputs if is_input else node.outputs
  417. counter = input_index if is_input else output_index
  418. reorder_collection = reorder_me_input if is_input else reorder_me_output
  419. if socket_map:
  420. if item.identifier in socket_map.keys():
  421. socket = relink_socket_map_add_socket(node, socket_collection, item, item.in_out)
  422. do_relink(node, socket, socket_map, item.in_out)
  423. else:
  424. for has_socket in socket_collection:
  425. if has_socket.bl_idname == item.bl_socket_idname and \
  426. has_socket.name == item.name:
  427. reorder_collection.append((has_socket, counter))
  428. break
  429. else:
  430. socket = relink_socket_map_add_socket(node, socket_collection, item, item.in_out)
  431. else:
  432. socket = relink_socket_map_add_socket(node, socket_collection, item, item.in_out)
  433. counter += 1
  434. # TODO: de-duplicate this hideous stuff
  435. for item in node.node_tree.interface.items_tree:
  436. if item.item_type != "SOCKET": continue
  437. if (item.in_out == 'INPUT' and update_input):
  438. # check and see if it exists... should only happen in curves on startup
  439. if item.bl_socket_idname in ['EnumCurveSocket']:
  440. for exists in node.inputs: # NOTE: check if the socket was not deleted
  441. if exists.identifier == item.identifier:
  442. # this happens for curve inputs because of some shennanigans with how
  443. # blender loads IDs - I can't set the ID until the file has loaded
  444. # so I have to avoid touching the socket until then...
  445. break
  446. else:
  447. update_group_sockets(item, True)
  448. else:
  449. update_group_sockets(item, True)
  450. input_index += 1
  451. if (item.in_out == 'OUTPUT' and update_output):
  452. if item.bl_socket_idname in ['EnumCurveSocket']: # LOOK up there at the comment!
  453. for exists in node.outputs:
  454. if exists.identifier == item.identifier:
  455. break
  456. else:
  457. update_group_sockets(item, False)
  458. else:
  459. update_group_sockets(item, False)
  460. output_index += 1
  461. both_reorders = zip([reorder_me_input, reorder_me_output], [node.inputs, node.outputs])
  462. for reorder_task, collection in both_reorders:
  463. for socket, position in reorder_task:
  464. for i, s in enumerate(collection): # get the index
  465. if s.identifier == socket.identifier: break
  466. else:
  467. prRed(f"WARN: could not reorder socket {socket.name}")
  468. to_index = position
  469. if (not socket.is_output) and node.bl_idname == "MantisSchemaGroup":
  470. to_index+=1
  471. collection.move(i, to_index)
  472. # at this point there is no wildcard socket
  473. if socket_map_in and '__extend__' in socket_map_in.keys():
  474. do_relink(node, None, socket_map_in, in_out='INPUT', parent_name='Constant' )
  475. node.id_data.do_live_update = toggle_update
  476. def node_tree_prop_update(self, context):
  477. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  478. return # so we check if an update is currently running.
  479. self.is_updating = True
  480. def init_schema(self, context):
  481. if len(self.inputs) == 0:
  482. self.inputs.new("UnsignedIntSocket", "Schema Length", identifier='Schema Length')
  483. if self.inputs[-1].bl_idname != "WildcardSocket":
  484. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  485. init_schema(self, context)
  486. try:
  487. node_group_update(self, force=True)
  488. finally: # ensure this line is run even if there is an error
  489. self.is_updating = False
  490. if self.bl_idname in ['MantisSchemaGroup'] and self.node_tree is not None:
  491. init_schema(self, context)
  492. from bpy.types import NodeCustomGroup
  493. def group_draw_buttons(self, context, layout):
  494. row = layout.row(align=True)
  495. row.prop(self, "node_tree", text="")
  496. if self.node_tree is None:
  497. row.operator("mantis.new_node_tree", text="", icon='PLUS', emboss=True)
  498. else:
  499. row.operator("mantis.edit_group", text="", icon='NODETREE', emboss=True)
  500. class MantisNodeGroup(Node, MantisUINode):
  501. bl_idname = "MantisNodeGroup"
  502. bl_label = "Node Group"
  503. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree, update=node_tree_prop_update,)
  504. is_updating:BoolProperty(default=False)
  505. def draw_label(self):
  506. if self.node_tree is None:
  507. return "Node Group"
  508. else:
  509. return self.node_tree.name
  510. def draw_buttons(self, context, layout):
  511. group_draw_buttons(self, context, layout)
  512. def update(self):
  513. if self.node_tree is None:
  514. return
  515. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  516. return # so we check if an update is currently running.
  517. live_update = self.id_data.do_live_update
  518. self.is_updating = True
  519. try:
  520. node_group_update(self)
  521. finally: # we need to reset this regardless of whether or not the operation succeeds!
  522. self.is_updating = False
  523. self.id_data.do_live_update = live_update # ensure this remains the same
  524. class GraphError(Exception):
  525. pass
  526. def get_signature_from_edited_tree(node, context):
  527. sig_path=[None,]
  528. for item in context.space_data.path[:-1]:
  529. sig_path.append(item.node_tree.nodes.active.name)
  530. return tuple(sig_path+[node.name])
  531. def poll_node_tree_schema(self, object):
  532. if isinstance(object, SchemaTree):
  533. return True
  534. return False
  535. # TODO tiny UI problem - inserting new links into the tree will not place them in the right place.
  536. class SchemaGroup(Node, MantisUINode):
  537. bl_idname = "MantisSchemaGroup"
  538. bl_label = "Node Schema"
  539. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree_schema, update=node_tree_prop_update,)
  540. is_updating:BoolProperty(default=False)
  541. def draw_buttons(self, context, layout):
  542. group_draw_buttons(self, context, layout)
  543. def draw_label(self):
  544. if self.node_tree is None:
  545. return "Schema Group"
  546. else:
  547. return self.node_tree.name
  548. def update(self):
  549. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  550. return # so we check if an update is currently running.
  551. if self.node_tree is None:
  552. return
  553. live_update = self.id_data.do_live_update
  554. self.is_updating = True
  555. try:
  556. node_group_update(self)
  557. # reset things if necessary:
  558. if self.node_tree:
  559. if len(self.inputs) == 0:
  560. self.inputs.new("UnsignedIntSocket", "Schema Length", identifier='Schema Length')
  561. if self.inputs[-1].identifier != "__extend__":
  562. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  563. finally: # we need to reset this regardless of whether or not the operation succeeds!
  564. self.is_updating = False
  565. self.id_data.do_live_update = live_update # ensure this remains the same
  566. # replace names with bl_idnames for reading the tree and solving schemas.
  567. replace_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",
  568. "SchemaArrayInput", "SchemaArrayInputAll", "SchemaConstInput", "SchemaConstOutput",
  569. "SchemaIndex", "SchemaOutgoingConnection", "SchemaArrayOutput","SchemaArrayInputGet",
  570. ]
  571. # anything that gets properties added in the graph... this is a clumsy approach but I need to watch for this
  572. # in schema generation and this is the easiest way to do it for now.
  573. custom_props_types = ["LinkArmature", "UtilityKeyframe", "UtilityFCurve", "UtilityDriver", "xFormBone"]
  574. # filters for determining if a link is a hierarchy link or a non-hierarchy (cyclic) link.
  575. from_name_filter = ["Driver",]
  576. to_name_filter = [
  577. "Custom Object xForm Override",
  578. "Custom Object",
  579. "Deform Bones",
  580. ]
  581. # nodes that must be solved as if they were Schema because they send arrays out.
  582. array_output_types = [
  583. 'UtilityArrayGet', 'UtilityKDChoosePoint', 'UtilityKDChooseXForm',
  584. ]
  585. def can_remove_socket_for_autogen(node, socket) -> bool:
  586. """ Whether to enable socket removal optimization for the socket
  587. This should be disallowed if e.g. it is a custom property.
  588. """
  589. return False # for now! This doesn't seem to be working...
  590. # the problem is that Schema does this, and so does Readtree
  591. # and they can try and both do it. That's bad.
  592. if node.socket_templates:
  593. for s_template in node.socket_templates:
  594. if s_template.name == socket:
  595. # raise NotImplementedError
  596. return True
  597. elif node.node_type == 'UTILITY':
  598. return True # HACK because most utilities don't have socket templates yet
  599. return False
  600. # TODO:
  601. # - get the execution context in the execution code
  602. # - from there, begin to use it for stuff I can't do without it
  603. # - and slowly start transferring stuff to it
  604. # The Mantis Overlay class is used to store node-tree specific information
  605. # such as inputs and outputs
  606. # used for e.g. allowing strings to pass as $variables in node names
  607. class MantisOverlay():
  608. def __init__( self, parent, inputs, outputs, ):
  609. pass
  610. # The MantisExecutionContext class is used to store the execution-specific variables
  611. # that are used when executing the tree
  612. # Importantly, it is NOT used to store variables between solutions, these belong to the
  613. # tree itself.
  614. class MantisExecutionContext():
  615. def __init__(
  616. self,
  617. base_tree,
  618. ):
  619. self.base_tree = base_tree
  620. self.execution_id = base_tree.execution_id
  621. self.execution_failed=False
  622. self.b_objects={} # objects created by Mantis during execution
  623. self.string_variables={}
  624. class MantisNode:
  625. """
  626. This class contains the basic interface for a Mantis Node.
  627. A MantisNode is used internally by Mantis to represent the final evaluated node graph.
  628. It gets generated with data from a MantisUINode when the graph is read.
  629. """
  630. def __init__(self, signature : tuple,
  631. base_tree : bpy.types.NodeTree,
  632. socket_templates : list[MantisSocketTemplate]=[],):
  633. self.base_tree=base_tree
  634. self.signature = signature
  635. self.ui_signature = signature
  636. self.inputs = MantisNodeSocketCollection(node=self, is_input=True)
  637. self.outputs = MantisNodeSocketCollection(node=self, is_input=False)
  638. self.parameters, self.drivers = {}, {}; self.bObject=None
  639. self.node_type='UNINITIALIZED'
  640. self.hierarchy_connections, self.connections = [], []
  641. self.hierarchy_dependencies, self.dependencies = [], []
  642. self.prepared, self.executed = False, False
  643. self.execution_prepared = False
  644. # the above is for tracking prep state in execution, so that I can avoid preparing nodes
  645. # again without changing the readtree code much.
  646. self.socket_templates = socket_templates
  647. self.mContext = None # for now I am gonna set this at runtime
  648. # I know it isn't "beautiful OOP" or whatever, but it is just easier
  649. # code should be simple and do things in the simplest way.
  650. # in the future I can refactor it, but it will require changes to 100+
  651. # classes, instead of adding about 5 lines of code elsewhere.
  652. if self.socket_templates:
  653. self.init_sockets()
  654. @property
  655. def name(self):
  656. return self.ui_signature[-1]
  657. @property
  658. def bl_idname(self): # this and the above exist solely to maintain interface w/bpy.types.Node
  659. from .utilities import get_node_prototype
  660. return get_node_prototype(self.ui_signature, self.base_tree).bl_idname
  661. def reset_execution(self) -> None:
  662. """ Reset the node for additional execution without re-building the tree."""
  663. self.drivers={}; self.bObject=None
  664. self.executed = False
  665. self.execution_prepared = False
  666. def init_sockets(self) -> None:
  667. self.inputs.init_sockets(self.socket_templates)
  668. self.outputs.init_sockets(self.socket_templates)
  669. def init_parameters(self, additional_parameters = {}) -> None:
  670. for socket in self.inputs:
  671. self.parameters[socket.name] = None
  672. for socket in self.outputs:
  673. self.parameters[socket.name] = None
  674. for key, value in additional_parameters.items():
  675. self.parameters[key]=value
  676. def gen_property_socket_map(self) -> dict:
  677. props_sockets = {}
  678. for s_template in self.socket_templates:
  679. if not s_template.blender_property:
  680. continue
  681. if isinstance(s_template.blender_property, str):
  682. props_sockets[s_template.blender_property]=(s_template.name, s_template.default_value)
  683. elif isinstance(s_template.blender_property, (tuple, list)):
  684. for index, sub_prop in enumerate(s_template.blender_property):
  685. props_sockets[sub_prop]=( (s_template.name, index),s_template.default_value[index] )
  686. return props_sockets
  687. def set_traverse(self, traversal_pairs = [(str, str)]) -> None:
  688. for (a, b) in traversal_pairs:
  689. self.inputs[a].set_traverse_target(self.outputs[b])
  690. self.outputs[b].set_traverse_target(self.inputs[a])
  691. def flush_links(self) -> None:
  692. for inp in self.inputs.values():
  693. inp.flush_links()
  694. for out in self.outputs.values():
  695. out.flush_links()
  696. def update_socket_value(self, blender_property, value) -> bool:
  697. change_handled=False
  698. if self.node_type == 'LINK':
  699. if len(self.bObject) == 0: # - there are no downstream xForms
  700. return True # so there is nothing to do here
  701. for b_ob in self.bObject:
  702. try:
  703. setattr(b_ob, blender_property, value)
  704. change_handled=True
  705. except Exception as e:
  706. print("Failed to update mantis socket because of %s" % e,
  707. "Updating tree instead.")
  708. else:
  709. try:
  710. b_ob = self.bObject
  711. if self.node_type == 'XFORM': # HACK
  712. b_ob = self.bGetObject()
  713. setattr(b_ob, blender_property, value)
  714. change_handled=True
  715. except Exception as e:
  716. print("Failed to update mantis socket because of %s" % e,
  717. "Updating tree instead.")
  718. return change_handled
  719. def ui_modify_socket(self, ui_socket, socket_name=None) -> bool:
  720. """ Handle changes in the node's UI. Updates the rig if possible."""
  721. # Always update the node's data
  722. change_handled=False
  723. if socket_name is None: socket_name = ui_socket.name
  724. value = ui_socket.default_value
  725. if socket_name == 'Enable': value = not value
  726. try:
  727. self.parameters[ui_socket.name]=value
  728. except KeyError:
  729. prRed(f"Unhandled change occured in socket {ui_socket.name} in node"
  730. f" {ui_socket.node.name} in tree {ui_socket.node.id_data.name}.")
  731. for s_template in self.socket_templates:
  732. if s_template.name==ui_socket.name:
  733. change_handled = True
  734. if not s_template.blender_property: return False
  735. elif isinstance(s_template.blender_property, str):
  736. change_handled &= self.update_socket_value(
  737. s_template.blender_property, value)
  738. else: # it is a tuple
  739. for i, prop in enumerate(s_template.blender_property):
  740. try:
  741. change_handled &= self.update_socket_value(
  742. prop, value[i])
  743. except IndexError:
  744. prRed(f"{ui_socket.name} does not have enough values to unpack"
  745. " to update the Mantis tree. Please report this as a bug.")
  746. change_handled=False
  747. break # we don't have to look through any more socket templates
  748. return change_handled
  749. # the goal here is to tag the node as unprepared
  750. # but some nodes are always prepared, so we have to kick it forward.
  751. def reset_execution_recursive(self):
  752. self.reset_execution()
  753. if self.prepared==False: return # all good from here
  754. for conn in self.hierarchy_connections:
  755. conn.reset_execution_recursive()
  756. def get_interface_signature(self):
  757. interface_signature = (*self.signature[:-1], 'InputInterface')
  758. ui_name = self.ui_signature[-1]
  759. name = self.signature[-1]
  760. if ui_name != name:
  761. schema_suffix = ''
  762. for i in range(len(name)):
  763. if name[i] == ui_name[i]: continue
  764. break
  765. schema_suffix = name[i+1:]
  766. interface_signature = (*interface_signature[:-1],
  767. 'InputInterface'+schema_suffix)
  768. return interface_signature
  769. # TODO: make this MUCH more efficient!
  770. # alternatively: call this ONCE when initializing the tree, precache results?
  771. def apply_string_variables(self, string):
  772. # We get the mContext, iterate through the signature, and string-replace variables
  773. # this function should be called by evaluate_input if the result is a string.
  774. result=string; name=""
  775. for i in range(len(self.signature[:-1])):
  776. if i == 0: continue # it is None or AUTOGENERATED or something
  777. name+=self.signature[i]
  778. vars = self.mContext.string_variables.get(name, None)
  779. if vars is None:
  780. raise RuntimeError("Can't get string variables for node")
  781. elif not vars:
  782. prWhite(f"INFO: No vars available for {self}")
  783. var_name_keys = list(vars.keys()); var_name_keys.sort(key=lambda a : -len(a))
  784. for var_name in var_name_keys:
  785. var_value = vars[var_name]
  786. if var_value is None:
  787. interface_signature = self.get_interface_signature()
  788. interface_node = self.base_tree.parsed_tree.get(interface_signature)
  789. from .utilities import set_string_variables_at_execution
  790. set_string_variables_at_execution(interface_node, var_name)
  791. var_value = vars[var_name]
  792. result = result.replace("$"+var_name, var_value)
  793. return result
  794. def evaluate_input(self, input_name, index=0) -> Any:
  795. from .node_container_common import trace_single_line
  796. if not (self.inputs.get(input_name)): # get the named parameter if there is no input
  797. return self.parameters.get(input_name) # this will return None if the parameter does not exist.
  798. # this trace() should give a key error if there is a problem
  799. # it is NOT handled here because it should NOT happen - so I want the error message.
  800. trace = trace_single_line(self, input_name, index)
  801. prop = trace[0][-1].parameters[trace[1].name] #trace[0] = the list of traced nodes; read its parameters
  802. # apply the string variables if possible
  803. if isinstance(prop, str) and "$" in prop: prop = self.apply_string_variables(prop)
  804. return prop
  805. def fill_parameters(self, ui_node=None) -> None:
  806. from .utilities import get_node_prototype
  807. from .node_container_common import get_socket_value
  808. if not ui_node:
  809. if ( (self.signature[0] in ["MANTIS_AUTOGENERATED", "SCHEMA_AUTOGENERATED" ]) or
  810. (self.signature[-1] in ["NodeGroupOutput", "NodeGroupInput"]) ): # I think this is harmless
  811. return None
  812. else: # BUG shouldn't this use ui_signature??
  813. ui_node = get_node_prototype(self.signature, self.base_tree)
  814. if not ui_node:
  815. raise RuntimeError(wrapRed("No node prototype found for... %s" % ( [self.base_tree] + list(self.signature[1:]) ) ) )
  816. for key in self.parameters.keys():
  817. node_socket = ui_node.inputs.get(key)
  818. if self.parameters[key] is not None: # the parameters are usually initialized as None.
  819. continue # will be filled by the node itself
  820. if not node_socket: #maybe the node socket has no name
  821. if ( ( len(ui_node.inputs) == 0) and ( len(ui_node.outputs) == 1) ):
  822. node_socket = ui_node.outputs[0] # this is a simple input node.
  823. elif key == 'Name': # for Links we just use the Node Label, or if there is no label, the name.
  824. self.parameters[key] = ui_node.label if ui_node.label else ui_node.name
  825. continue
  826. if node_socket:
  827. if node_socket.bl_idname in ['RelationshipSocket', 'xFormSocket']: continue
  828. elif node_socket.is_linked and (not node_socket.is_output): continue
  829. # we will get the value from the link, because this is a linked input port.
  830. # very importantly, we do not pass linked outputs
  831. # fill these because they are probably Input nodes.
  832. elif hasattr(node_socket, "default_value"):
  833. if (value := get_socket_value(node_socket)) is not None:
  834. self.parameters[key] = value
  835. else:
  836. raise RuntimeError(wrapRed("No value found for " + self.__repr__() + " when filling out node parameters for " + ui_node.name + "::"+node_socket.name))
  837. # I don't think this works! but I like the idea
  838. def call_on_all_ancestors(self, *args, **kwargs):
  839. """Resolve the dependencies of this node with the named method and its arguments.
  840. First, dependencies are discovered by walking backwards through the tree. Once the root
  841. nodes are discovered, the method is called by each node in dependency order.
  842. The first argument MUST be the name of the method as a string.
  843. """
  844. if args[0] == 'call_on_all_ancestors': raise RuntimeError("Very funny!")
  845. from .utilities import get_all_dependencies
  846. from collections import deque
  847. # get all dependencies by walking backward through the tree.
  848. all_dependencies = get_all_dependencies(self)
  849. # get just the roots
  850. can_solve = deque(filter(lambda a : len(a.hierarchy_connections) == 0, all_dependencies))
  851. solved = set()
  852. while can_solve:
  853. node = can_solve.pop()
  854. method = getattr(node, args[0])
  855. method(*args[0:], **kwargs)
  856. solved.add(node)
  857. can_solve.extendleft(filter(lambda a : a in all_dependencies, node.hierarchy_connections))
  858. if self in solved:
  859. break
  860. return
  861. # gets targets for constraints and deformers and should handle all cases
  862. def get_target_and_subtarget(self, constraint_or_deformer, input_name = "Target"):
  863. from bpy.types import PoseBone, Object, SplineIKConstraint
  864. subtarget = ''; target = self.evaluate_input(input_name)
  865. if target:
  866. if not hasattr(target, "bGetObject"):
  867. if hasattr(constraint_or_deformer, 'name'):
  868. name = constraint_or_deformer.name
  869. else:
  870. name = 'NAME NOT FOUND'
  871. prRed(f"No {input_name} target found for {name} in {self} because there is no connected node, or node is wrong type")
  872. return
  873. if (isinstance(target.bGetObject(), PoseBone)):
  874. subtarget = target.bGetObject().name
  875. target = target.bGetParentArmature()
  876. elif (isinstance(target.bGetObject(), Object) ):
  877. target = target.bGetObject()
  878. else:
  879. raise RuntimeError("Cannot interpret constraint or deformer target!")
  880. if (isinstance(constraint_or_deformer, SplineIKConstraint)):
  881. if target and target.type not in ["CURVE"]:
  882. raise GraphError(wrapRed("Error: %s requires a Curve input, not %s" %
  883. (self, type(target))))
  884. constraint_or_deformer.target = target# don't get a subtarget
  885. if (input_name == 'Pole Target'):
  886. constraint_or_deformer.pole_target, constraint_or_deformer.pole_subtarget = target, subtarget
  887. else:
  888. if hasattr(constraint_or_deformer, "target"):
  889. constraint_or_deformer.target = target
  890. if hasattr(constraint_or_deformer, "object"):
  891. constraint_or_deformer.object = target
  892. if hasattr(constraint_or_deformer, "subtarget"):
  893. constraint_or_deformer.subtarget = subtarget
  894. # PASSES DEFINED HERE!
  895. def bPrepare(self, bContext=None):
  896. return # This one runs BEFORE anything else
  897. def bTransformPass(self, bContext=None):
  898. return # This one runs in EDIT MODE
  899. def bRelationshipPass(self, bContext=None):
  900. return # This one runs in POSE MODE
  901. def bFinalize(self, bContext=None):
  902. return
  903. def bModifierApply(self, bContext=None):
  904. return
  905. if environ.get("DOERROR"):
  906. def __repr__(self):
  907. return self.signature.__repr__()
  908. else:
  909. def __repr__(self):
  910. return self.ui_signature.__repr__()
  911. # do I need this and the link class above?
  912. class DummyLink:
  913. #gonna use this for faking links to keep the interface consistent
  914. def __init__(self, from_socket, to_socket, nc_from=None, nc_to=None, original_from=None, multi_input_sort_id=0):
  915. self.from_socket = from_socket
  916. self.to_socket = to_socket
  917. self.nc_from = nc_from
  918. self.nc_to = nc_to
  919. self.multi_input_sort_id = multi_input_sort_id
  920. # self.from_node = from_socket.node
  921. # self.to_node = to_socket.node
  922. if (original_from):
  923. self.original_from = original_from
  924. else:
  925. self.original_from = self.from_socket
  926. def __repr__(self):
  927. return(self.nc_from.__repr__()+":"+self.from_socket.name + " -> " + self.nc_to.__repr__()+":"+self.to_socket.name)
  928. def detect_hierarchy_link(from_node, from_socket, to_node, to_socket,):
  929. if to_node.node_type in ['DUMMY_SCHEMA', 'SCHEMA']:
  930. return False #TODO: find out if filtering SCHEMA types is wise
  931. if (from_socket in from_name_filter) or (to_socket in to_name_filter):
  932. return False
  933. # if from_node.__class__.__name__ in ["UtilityCombineVector", "UtilityCombineThreeBool"]:
  934. # return False
  935. return True
  936. class NodeLink:
  937. from_node = None
  938. from_socket = None
  939. to_node = None
  940. to_socket = None
  941. def __init__(self, from_node, from_socket, to_node, to_socket, multi_input_sort_id=0, sub_sort_id=0):
  942. if from_node.signature == to_node.signature:
  943. raise RuntimeError("Cannot connect a node to itself.")
  944. self.from_node = from_node
  945. self.from_socket = from_socket
  946. self.to_node = to_node
  947. self.to_socket = to_socket
  948. self.from_node.outputs[self.from_socket].links.append(self)
  949. # it is the responsibility of the node that uses these links to sort them correctly based on the sort_id
  950. self.multi_input_sort_id = multi_input_sort_id # this is the sort_id of the link in the UI
  951. self.sub_sort_id = sub_sort_id # this is for sorting within a bundled link (one link in the UI)
  952. self.to_node.inputs[self.to_socket].links.append(self)
  953. self.is_hierarchy = detect_hierarchy_link(from_node, from_socket, to_node, to_socket,)
  954. self.is_alive = True
  955. def __repr__(self):
  956. return self.from_node.outputs[self.from_socket].__repr__() + " --> " + self.to_node.inputs[self.to_socket].__repr__()
  957. # link_string = # if I need to colorize output for debugging.
  958. # if self.is_hierarchy:
  959. # return wrapOrange(link_string)
  960. # else:
  961. # return wrapWhite(link_string)
  962. def die(self):
  963. self.is_alive = False
  964. self.to_node.inputs[self.to_socket].flush_links()
  965. self.from_node.outputs[self.from_socket].flush_links()
  966. def insert_node(self, middle_node, middle_node_in, middle_node_out, re_init_hierarchy = True):
  967. to_node = self.to_node
  968. to_socket = self.to_socket
  969. self.to_node = middle_node
  970. self.to_socket = middle_node_in
  971. middle_node.outputs[middle_node_out].connect(to_node, to_socket)
  972. if re_init_hierarchy:
  973. from .utilities import init_connections, init_dependencies
  974. init_connections(self.from_node)
  975. init_connections(middle_node)
  976. init_dependencies(middle_node)
  977. init_dependencies(to_node)
  978. class NodeSocket:
  979. # @property # this is a read-only property.
  980. # def is_linked(self):
  981. # return bool(self.links)
  982. def __init__(self, is_input = False,
  983. node = None, name = None,
  984. traverse_target = None):
  985. self.can_traverse = False # to/from the other side of the parent node
  986. self.traverse_target = None
  987. self.node = node
  988. self.name = name
  989. self.is_input = is_input
  990. self.links = []
  991. self.is_linked = False
  992. if (traverse_target):
  993. self.can_traverse = True
  994. def connect(self, node, socket, sort_id=0, sub_sort_id=0):
  995. if (self.is_input):
  996. to_node = self.node; from_node = node
  997. to_socket = self.name; from_socket = socket
  998. else:
  999. from_node = self.node; to_node = node
  1000. from_socket = self.name; to_socket = socket
  1001. from_node.outputs[from_socket].is_linked = True
  1002. to_node.inputs[to_socket].is_linked = True
  1003. # NOTE: I have removed a check for duplicate links here.
  1004. # Schemas sometimes have valid duplicate links.
  1005. # It is conceivable that this will lead to bugs, but I judge it unlikely.
  1006. new_link = NodeLink(
  1007. from_node,
  1008. from_socket,
  1009. to_node,
  1010. to_socket,
  1011. sort_id,
  1012. sub_sort_id)
  1013. return new_link
  1014. def set_traverse_target(self, traverse_target):
  1015. self.traverse_target = traverse_target
  1016. self.can_traverse = True
  1017. def flush_links(self):
  1018. """ Removes dead links from this socket."""
  1019. self.links = [l for l in self.links if l.is_alive]
  1020. self.links.sort(key=links_sort_key)
  1021. self.is_linked = bool(self.links)
  1022. @property
  1023. def is_connected(self):
  1024. return len(self.links)>0
  1025. def __repr__(self):
  1026. return self.node.__repr__() + "::" + self.name
  1027. class MantisNodeSocketCollection(dict):
  1028. def __init__(self, node, is_input=False):
  1029. self.is_input = is_input
  1030. self.node = node
  1031. def init_sockets(self, sockets):
  1032. for socket in sockets:
  1033. if isinstance(socket, str):
  1034. self[socket] = NodeSocket(is_input=self.is_input, name=socket, node=self.node)
  1035. elif isinstance(socket, MantisSocketTemplate):
  1036. if socket.is_input != self.is_input: continue
  1037. self[socket.name] = NodeSocket(is_input=self.is_input, name=socket.name, node=self.node)
  1038. else:
  1039. raise RuntimeError(f"NodeSocketCollection keys must be str or MantisSocketTemplate, not {type(socket)}")
  1040. def __delitem__(self, key):
  1041. """Deletes a node socket by name, and all its links."""
  1042. socket = self[key]
  1043. for l in socket.links:
  1044. l.die()
  1045. super().__delitem__(key)
  1046. def __iter__(self):
  1047. """Makes the class iterable"""
  1048. return iter(self.values())