base_definitions.py 42 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977
  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. def TellClasses():
  14. #Why use a function to do this? Because I don't need every class to register.
  15. return [ MantisTree,
  16. SchemaTree,
  17. MantisNodeGroup,
  18. SchemaGroup,
  19. ]
  20. def error_popup_draw(self, context):
  21. self.layout.label(text="Error executing tree. See Console.")
  22. mantis_root = ".".join(__name__.split('.')[:-1]) # absolute HACK
  23. # https://docs.blender.org/api/master/bpy.types.NodeTree.html#bpy.types.NodeTree.valid_socket_type
  24. # thank you, Sverchok
  25. def valid_interface_types(cls : NodeTree, socket_idname : str):
  26. from .socket_definitions import tell_valid_bl_idnames, TellClasses
  27. #TODO: do the versioning code to handle this so it can be in all versions
  28. if bpy.app.version <= (4,4,0): # should work in 4.4.1
  29. return socket_idname in [cls.bl_idname for cls in TellClasses()]
  30. else: # once versioning is finished this will be unnecesary.
  31. return socket_idname in tell_valid_bl_idnames()
  32. def fix_reroute_colors(tree):
  33. context = bpy.context
  34. if any((tree.is_executing, tree.is_exporting, tree.do_live_update==False, context.space_data is None) ):
  35. return
  36. from collections import deque
  37. from .utilities import socket_seek
  38. from .socket_definitions import MantisSocket
  39. reroutes_without_color = deque()
  40. for n in tree.nodes:
  41. if n.bl_idname=='NodeReroute' and n.inputs[0].bl_idname == "NodeSocketColor":
  42. reroutes_without_color.append(n)
  43. try:
  44. while reroutes_without_color:
  45. rr = reroutes_without_color.pop()
  46. if rr.inputs[0].is_linked:
  47. link = rr.inputs[0].links[0]
  48. from_node = link.from_node
  49. socket = socket_seek(link, tree.links)
  50. if isinstance(socket, MantisSocket):
  51. rr.socket_idname = socket.bl_idname
  52. except Exception as e:
  53. print(wrapOrange("WARN: Updating reroute color failed with exception: ")+wrapWhite(f"{e.__class__.__name__}"))
  54. class MantisTree(NodeTree):
  55. '''A custom node tree type that will show up in the editor type list'''
  56. bl_idname = 'MantisTree'
  57. bl_label = "Rigging Nodes"
  58. bl_icon = 'OUTLINER_OB_ARMATURE'
  59. tree_valid:BoolProperty(default=False)
  60. do_live_update:BoolProperty(default=True) # use this to disable updates for e.g. scripts
  61. num_links:IntProperty(default=-1)
  62. filepath:StringProperty(default="", subtype='FILE_PATH')
  63. is_executing:BoolProperty(default=False)
  64. is_exporting:BoolProperty(default=False)
  65. execution_id:StringProperty(default='')
  66. mantis_version:IntVectorProperty(default=[0,9,2])
  67. # this prevents the node group from executing on the next depsgraph update
  68. # because I don't always have control over when the dg update happens.
  69. prevent_next_exec:BoolProperty(default=False)
  70. parsed_tree={}
  71. if (bpy.app.version < (4, 4, 0)): # in 4.4 this leads to a crash
  72. @classmethod
  73. def valid_socket_type(cls : NodeTree, socket_idname: str):
  74. return valid_interface_types(cls, socket_idname)
  75. def update(self): # set the reroute colors
  76. if (bpy.app.version >= (4,4,0)):
  77. fix_reroute_colors(self)
  78. def update_tree(self, context = None):
  79. if self.is_exporting:
  80. return
  81. # return
  82. self.is_executing = True
  83. from . import readtree
  84. prGreen("Validating Tree: %s" % self.name)
  85. try:
  86. self.parsed_tree = readtree.parse_tree(self)
  87. if context:
  88. self.display_update(context)
  89. self.is_executing = False
  90. self.tree_valid = True
  91. except GraphError as e:
  92. prRed("Failed to update node tree due to error.")
  93. self.tree_valid = False
  94. self.is_executing = False
  95. raise e
  96. finally:
  97. self.is_executing = False
  98. def display_update(self, context):
  99. if self.is_exporting:
  100. return
  101. self.is_executing = True
  102. current_tree = bpy.context.space_data.path[-1].node_tree
  103. for node in current_tree.nodes:
  104. if hasattr(node, "display_update"):
  105. try:
  106. node.display_update(self.parsed_tree, context)
  107. except Exception as e:
  108. print("Node \"%s\" failed to update display with error: %s" %(wrapGreen(node.name), wrapRed(e)))
  109. self.is_executing = False
  110. # TODO: deal with invalid links properly.
  111. # - Non-hierarchy links should be ignored in the circle-check and so the links should be marked valid in such a circle
  112. # - hierarchy-links should be marked invalid and prevent the tree from executing.
  113. def execute_tree(self,context, error_popups = False):
  114. self.prevent_next_exec = False
  115. if self.is_exporting:
  116. return
  117. # return
  118. prGreen("Executing Tree: %s" % self.name)
  119. self.is_executing = True
  120. from . import readtree
  121. try:
  122. readtree.execute_tree(self.parsed_tree, self, context, error_popups)
  123. except RecursionError as e:
  124. prRed("Recursion error while parsing tree.")
  125. finally:
  126. self.is_executing = False
  127. class SchemaTree(NodeTree):
  128. '''A node tree representing a schema to generate a Mantis tree'''
  129. bl_idname = 'SchemaTree'
  130. bl_label = "Rigging Nodes Schema"
  131. bl_icon = 'RIGID_BODY_CONSTRAINT'
  132. # these are only needed for consistent interface, but should not be used
  133. do_live_update:BoolProperty(default=True) # default to true so that updates work
  134. is_executing:BoolProperty(default=False)
  135. is_exporting:BoolProperty(default=False)
  136. mantis_version:IntVectorProperty(default=[0,9,2])
  137. if (bpy.app.version < (4, 4, 0)): # in 4.4 this leads to a crash
  138. @classmethod
  139. def valid_socket_type(cls : NodeTree, socket_idname: str):
  140. return valid_interface_types(cls, socket_idname)
  141. def update(self): # set the reroute colors
  142. if (bpy.app.version >= (4,4,0)):
  143. fix_reroute_colors(self)
  144. from dataclasses import dataclass, field
  145. from typing import Any
  146. @dataclass
  147. class MantisSocketTemplate():
  148. name : str = field(default="")
  149. bl_idname : str = field(default="")
  150. traverse_target : str = field(default="")
  151. identifier : str = field(default="")
  152. display_shape : str = field(default="") # for arrays
  153. category : str = field(default="") # for use in display update
  154. blender_property : str | tuple[str] = field(default="") # for props_sockets -> evaluate sockets
  155. is_input : bool = field(default=False)
  156. hide : bool = field(default=False)
  157. use_multi_input : bool = field(default=False)
  158. default_value : Any = field(default=None)
  159. #TODO: do a better job explaining how MantisNode and MantisUINode relate.
  160. class MantisUINode:
  161. """
  162. This class contains the common user-interface features of Mantis nodes.
  163. MantisUINode objects will spawn one or more MantisNode objects when the graph is evaluated.
  164. The MantisNode objects will pull the data from the UI node and use it to generate the graph.
  165. """
  166. mantis_node_library=''
  167. mantis_node_class_name=''
  168. mantis_class=None
  169. @classmethod
  170. def poll(cls, ntree):
  171. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  172. @classmethod
  173. def set_mantis_class(self):
  174. from importlib import import_module
  175. # do not catch errors, they should cause a failure.
  176. try:
  177. module = import_module(self.mantis_node_library, package=mantis_root)
  178. self.mantis_class=getattr(module, self.mantis_node_class_name)
  179. except Exception as e:
  180. print(self)
  181. raise e
  182. def insert_link(self, link):
  183. if (bpy.app.version >= (4, 4, 0)):
  184. return # this causes a crash due to a bug.
  185. context = bpy.context
  186. if context.space_data:
  187. node_tree = context.space_data.path[0].node_tree
  188. if node_tree.do_live_update:
  189. node_tree.update_tree(context)
  190. if (link.to_socket.is_linked == False):
  191. node_tree.num_links+=1
  192. elif (link.to_socket.is_multi_input):
  193. node_tree.num_links+=1
  194. def init_sockets(self, socket_templates : tuple[MantisSocketTemplate]):
  195. for template in socket_templates:
  196. collection = self.outputs
  197. if template.is_input:
  198. collection = self.inputs
  199. identifier = template.name
  200. if template.identifier:
  201. identifier = template.identifier
  202. socket = collection.new(
  203. template.bl_idname,
  204. template.name,
  205. identifier=identifier,
  206. use_multi_input=template.use_multi_input
  207. )
  208. socket.hide= template.hide
  209. if template.category:
  210. # a custom property for the UI functions to use.
  211. socket['category'] = template.category
  212. if template.default_value is not None:
  213. socket.default_value = template.default_value
  214. # this can throw a TypeError - it is the caller's
  215. # responsibility to send the right type.
  216. class SchemaUINode(MantisUINode):
  217. mantis_node_library='.schema_containers'
  218. @classmethod
  219. def poll(cls, ntree):
  220. return (ntree.bl_idname in ['SchemaTree'])
  221. class LinkNode(MantisUINode):
  222. mantis_node_library='.link_containers'
  223. @classmethod
  224. def poll(cls, ntree):
  225. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  226. class xFormNode(MantisUINode):
  227. mantis_node_library='.xForm_containers'
  228. @classmethod
  229. def poll(cls, ntree):
  230. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  231. class DeformerNode(MantisUINode):
  232. mantis_node_library='.deformer_containers'
  233. @classmethod
  234. def poll(cls, ntree):
  235. return (ntree.bl_idname in ['MantisTree', 'SchemaTree'])
  236. def poll_node_tree(self, object):
  237. forbid = []
  238. context = bpy.context
  239. if context.space_data:
  240. if context.space_data.path:
  241. for path_item in context.space_data.path:
  242. forbid.append(path_item.node_tree.name)
  243. if isinstance(object, MantisTree) and object.name not in forbid:
  244. return True
  245. return False
  246. # TODO: try and remove the extra loop used here... but it is OK for now
  247. def should_remove_socket(node, socket):
  248. # a function to check if the socket is in the interface
  249. id_found = False
  250. for item in node.node_tree.interface.items_tree:
  251. if item.item_type != "SOCKET": continue
  252. if item.identifier == socket.identifier:
  253. id_found = True; break
  254. return not id_found
  255. # TODO: try to check identifiers instead of name.
  256. def node_group_update(node, force = False):
  257. if not node.is_updating:
  258. raise RuntimeError("Cannot update node while it is not marked as updating.")
  259. if not force:
  260. if (node.id_data.do_live_update == False) or \
  261. (node.id_data.is_executing == True) or \
  262. (node.id_data.is_exporting == True):
  263. return
  264. # note: if (node.id_data.is_exporting == True) I need to be able to update so I can make links.
  265. toggle_update = node.id_data.do_live_update
  266. node.id_data.do_live_update = False
  267. identifiers_in={socket.identifier:socket for socket in node.inputs}
  268. identifiers_out={socket.identifier:socket for socket in node.outputs}
  269. indices_in,indices_out={},{} # check by INDEX to see if the socket's name/type match.
  270. for collection, map in [(node.inputs, indices_in), (node.outputs, indices_out)]:
  271. for i, socket in enumerate(collection):
  272. map[socket.identifier]=i
  273. if node.node_tree is None:
  274. node.inputs.clear(); node.outputs.clear()
  275. node.id_data.do_live_update = toggle_update
  276. return
  277. found_in, found_out = [], []
  278. update_input, update_output = False, False
  279. for item in node.node_tree.interface.items_tree:
  280. if item.item_type != "SOCKET": continue
  281. if item.in_out == 'OUTPUT':
  282. if s:= identifiers_out.get(item.identifier): # if the requested output doesn't exist, update
  283. found_out.append(item.identifier)
  284. if (indices_out[s.identifier]!=item.index): update_output=True; continue
  285. if update_output: continue
  286. if s.bl_idname != item.socket_type: update_output = True; continue
  287. else: update_output = True; continue
  288. else:
  289. if s:= identifiers_in.get(item.identifier): # if the requested input doesn't exist, update
  290. found_in.append(item.identifier)
  291. if (indices_in[s.identifier]!=item.index): update_input=True; continue
  292. if update_input: continue # done here
  293. if s.bl_idname != item.socket_type: update_input = True; continue
  294. else: update_input = True; continue
  295. # Schema has an extra input for Length and for Extend.
  296. if node.bl_idname == 'MantisSchemaGroup':
  297. found_in.extend(['Schema Length', ''])
  298. # if we have too many elements, just get rid of the ones we don't need
  299. if len(node.inputs) > len(found_in):#
  300. for inp in node.inputs:
  301. if inp.identifier in found_in: continue
  302. node.inputs.remove(inp)
  303. if len(node.outputs) > len(found_out):
  304. for out in node.outputs:
  305. if out.identifier in found_out: continue
  306. node.outputs.remove(out)
  307. #
  308. if len(node.inputs) > 0 and (inp := node.inputs[-1]).bl_idname == 'WildcardSocket' and inp.is_linked:
  309. update_input = True
  310. #
  311. if not (update_input or update_output):
  312. node.id_data.do_live_update = toggle_update
  313. return
  314. if update_input or update_output:
  315. socket_maps = get_socket_maps(node,)
  316. if socket_maps:
  317. socket_map_in, socket_map_out = socket_maps
  318. if node.bl_idname == "MantisSchemaGroup" and \
  319. len(node.inputs)+len(node.outputs)<=2 and\
  320. len(node.node_tree.interface.items_tree) > 0:
  321. socket_map_in, socket_map_out = None, None
  322. # We have to initialize the node because it only has its base inputs.
  323. elif socket_maps is None:
  324. node.id_data.do_live_update = toggle_update
  325. return
  326. if update_input :
  327. if node.bl_idname == 'MantisSchemaGroup':
  328. schema_length=0
  329. if sl := node.inputs.get("Schema Length"):
  330. schema_length = sl.default_value
  331. # sometimes this isn't available yet # TODO not happy about this solution
  332. remove_me=[]
  333. # remove all found map items but the Schema Length input (reuse it)
  334. for i, socket in enumerate(node.inputs):
  335. if socket.identifier == "Schema Length" and i == 0:
  336. continue
  337. elif (socket_map_in is None) or socket.identifier in socket_map_in.keys():
  338. remove_me.append(socket)
  339. elif should_remove_socket(node, socket):
  340. remove_me.append(socket)
  341. while remove_me:
  342. node.inputs.remove(remove_me.pop())
  343. if update_output:
  344. remove_me=[]
  345. for socket in node.outputs:
  346. if (socket_map_out is None) or socket.identifier in socket_map_out.keys():
  347. remove_me.append(socket)
  348. elif should_remove_socket(node, socket):
  349. remove_me.append(socket)
  350. while remove_me:
  351. node.inputs.remove(remove_me.pop())
  352. from .utilities import relink_socket_map_add_socket
  353. reorder_me_input = []; input_index = 0
  354. reorder_me_output = []; output_index = 0
  355. def update_group_sockets(interface_item, is_input):
  356. socket_map = socket_map_in if is_input else socket_map_out
  357. socket_collection = node.inputs if is_input else node.outputs
  358. counter = input_index if is_input else output_index
  359. reorder_collection = reorder_me_input if is_input else reorder_me_output
  360. if socket_map:
  361. if item.identifier in socket_map.keys():
  362. socket = relink_socket_map_add_socket(node, socket_collection, item)
  363. do_relink(node, socket, socket_map, item.in_out)
  364. else:
  365. for has_socket in socket_collection:
  366. if has_socket.bl_idname == item.socket_type and \
  367. has_socket.name == item.name:
  368. reorder_collection.append((has_socket, counter))
  369. break
  370. else:
  371. socket = relink_socket_map_add_socket(node, socket_collection, item)
  372. else:
  373. socket = relink_socket_map_add_socket(node, socket_collection, item)
  374. counter += 1
  375. for item in node.node_tree.interface.items_tree:
  376. if item.item_type != "SOCKET": continue
  377. if (item.in_out == 'INPUT' and update_input):
  378. update_group_sockets(item, True)
  379. if (item.in_out == 'OUTPUT' and update_output):
  380. update_group_sockets(item, False)
  381. both_reorders = zip([reorder_me_input, reorder_me_output], [node.inputs, node.outputs])
  382. for reorder_task, collection in both_reorders:
  383. for socket, position in reorder_task:
  384. for i, s in enumerate(collection): # get the index
  385. if s.identifier == socket.identifier: break
  386. else:
  387. prRed(f"WARN: could not reorder socket {socket.name}")
  388. to_index = position
  389. if (not socket.is_output) and node.bl_idname == "MantisSchemaGroup":
  390. to_index+=1
  391. collection.move(i, to_index)
  392. # at this point there is no wildcard socket
  393. if socket_map_in and '__extend__' in socket_map_in.keys():
  394. do_relink(node, None, socket_map_in, in_out='INPUT', parent_name='Constant' )
  395. node.id_data.do_live_update = toggle_update
  396. def node_tree_prop_update(self, context):
  397. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  398. return # so we check if an update is currently running.
  399. self.is_updating = True
  400. def init_schema(self, context):
  401. if len(self.inputs) == 0:
  402. self.inputs.new("UnsignedIntSocket", "Schema Length", identifier='Schema Length')
  403. if self.inputs[-1].bl_idname != "WildcardSocket":
  404. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  405. init_schema(self, context)
  406. try:
  407. node_group_update(self, force=True)
  408. finally: # ensure this line is run even if there is an error
  409. self.is_updating = False
  410. if self.bl_idname in ['MantisSchemaGroup'] and self.node_tree is not None:
  411. init_schema(self, context)
  412. from bpy.types import NodeCustomGroup
  413. def group_draw_buttons(self, context, layout):
  414. row = layout.row(align=True)
  415. row.prop(self, "node_tree", text="")
  416. if self.node_tree is None:
  417. row.operator("mantis.new_node_tree", text="", icon='PLUS', emboss=True)
  418. else:
  419. row.operator("mantis.edit_group", text="", icon='NODETREE', emboss=True)
  420. class MantisNodeGroup(Node, MantisUINode):
  421. bl_idname = "MantisNodeGroup"
  422. bl_label = "Node Group"
  423. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree, update=node_tree_prop_update,)
  424. is_updating:BoolProperty(default=False)
  425. def draw_label(self):
  426. if self.node_tree is None:
  427. return "Node Group"
  428. else:
  429. return self.node_tree.name
  430. def draw_buttons(self, context, layout):
  431. group_draw_buttons(self, context, layout)
  432. def update(self):
  433. if self.node_tree is None:
  434. return
  435. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  436. return # so we check if an update is currently running.
  437. live_update = self.id_data.do_live_update
  438. self.is_updating = True
  439. try:
  440. node_group_update(self)
  441. finally: # we need to reset this regardless of whether or not the operation succeeds!
  442. self.is_updating = False
  443. self.id_data.do_live_update = live_update # ensure this remains the same
  444. class GraphError(Exception):
  445. pass
  446. def get_signature_from_edited_tree(node, context):
  447. sig_path=[None,]
  448. for item in context.space_data.path[:-1]:
  449. sig_path.append(item.node_tree.nodes.active.name)
  450. return tuple(sig_path+[node.name])
  451. def poll_node_tree_schema(self, object):
  452. if isinstance(object, SchemaTree):
  453. return True
  454. return False
  455. # TODO tiny UI problem - inserting new links into the tree will not place them in the right place.
  456. class SchemaGroup(Node, MantisUINode):
  457. bl_idname = "MantisSchemaGroup"
  458. bl_label = "Node Schema"
  459. node_tree:PointerProperty(type=NodeTree, poll=poll_node_tree_schema, update=node_tree_prop_update,)
  460. is_updating:BoolProperty(default=False)
  461. def draw_buttons(self, context, layout):
  462. group_draw_buttons(self, context, layout)
  463. def draw_label(self):
  464. if self.node_tree is None:
  465. return "Schema Group"
  466. else:
  467. return self.node_tree.name
  468. def update(self):
  469. if self.is_updating: # update() can be called from update() and that leads to an infinite loop.
  470. return # so we check if an update is currently running.
  471. if self.node_tree is None:
  472. return
  473. live_update = self.id_data.do_live_update
  474. self.is_updating = True
  475. try:
  476. node_group_update(self)
  477. # reset things if necessary:
  478. if self.node_tree:
  479. if len(self.inputs) == 0:
  480. self.inputs.new("UnsignedIntSocket", "Schema Length", identifier='Schema Length')
  481. if self.inputs[-1].identifier != "__extend__":
  482. self.inputs.new("WildcardSocket", "", identifier="__extend__")
  483. finally: # we need to reset this regardless of whether or not the operation succeeds!
  484. self.is_updating = False
  485. self.id_data.do_live_update = live_update # ensure this remains the same
  486. NODES_REMOVED=["xFormRootNode"]
  487. # Node bl_idname, # Socket Name
  488. SOCKETS_REMOVED=[("UtilityDriverVariable", "Transform Channel"),
  489. ("xFormRootNode","World Out"),
  490. ("UtilitySwitch","xForm"),
  491. ("LinkDrivenParameter", "Enable")]
  492. # Node Class #Prior bl_idname # prior name # new bl_idname # new name, # Multi
  493. SOCKETS_RENAMED=[ ("LinkDrivenParameter", "DriverSocket", "Driver", "FloatSocket", "Value", False),
  494. ("DeformerHook", "IntSocket", "Index", "UnsignedIntSocket", "Curve Point Index", False)]
  495. # NODE CLASS NAME IN_OUT SOCKET TYPE SOCKET NAME INDEX MULTI DEFAULT
  496. SOCKETS_ADDED=[("DeformerMorphTargetDeform", 'INPUT', 'BooleanSocket', "Use Shape Key", 1, False, False),
  497. ("DeformerMorphTargetDeform", 'INPUT', 'BooleanSocket', "Use Offset", 2, False, True),
  498. ("UtilityFCurve", 'INPUT', "eFCrvExtrapolationMode", "Extrapolation Mode", 0, False, 'CONSTANT'),
  499. ("LinkCopyScale", 'INPUT', "BooleanSocket", "Additive", 3, False, False),
  500. ("DeformerHook", 'INPUT', "FloatFactorSocket", "Influence", 3, False, 1.0),
  501. ("DeformerHook", 'INPUT', "BooleanSocket", "Auto-Bezier", 4, False, True),
  502. ("UtilityCompare", 'INPUT', "EnumCompareOperation", "Comparison", 0, False, 'EQUAL'),
  503. ]
  504. # replace names with bl_idnames for reading the tree and solving schemas.
  505. replace_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",
  506. "SchemaArrayInput", "SchemaArrayInputAll", "SchemaConstInput", "SchemaConstOutput",
  507. "SchemaIndex", "SchemaOutgoingConnection", "SchemaArrayOutput","SchemaArrayInputGet",
  508. ]
  509. # anything that gets properties added in the graph... this is a clumsy approach but I need to watch for this
  510. # in schema generation and this is the easiest way to do it for now.
  511. custom_props_types = ["LinkArmature", "UtilityKeyframe", "UtilityFCurve", "UtilityDriver", "xFormBone"]
  512. # filters for determining if a link is a hierarchy link or a non-hierarchy (cyclic) link.
  513. from_name_filter = ["Driver",]
  514. to_name_filter = [
  515. "Custom Object xForm Override",
  516. "Custom Object",
  517. "Deform Bones",
  518. ]
  519. # nodes that must be solved as if they were Schema because they send arrays out.
  520. array_output_types = [
  521. 'UtilityArrayGet', 'UtilityKDChoosePoint', 'UtilityKDChooseXForm',
  522. ]
  523. # TODO:
  524. # - get the execution context in the execution code
  525. # - from there, begin to use it for stuff I can't do without it
  526. # - and slowly start transferring stuff to it
  527. # The Mantis Overlay class is used to store node-tree specific information
  528. # such as inputs and outputs
  529. # used for e.g. allowing strings to pass as $variables in node names
  530. class MantisOverlay():
  531. def __init__( self, parent, inputs, outputs, ):
  532. pass
  533. # The MantisExecutionContext class is used to store the execution-specific variables
  534. # that are used when executing the tree
  535. # Importantly, it is NOT used to store variables between solutions, these belong to the
  536. # tree itself.
  537. class MantisExecutionContext():
  538. def __init__(
  539. self,
  540. base_tree,
  541. ):
  542. self.base_tree = base_tree
  543. self.execution_id = base_tree.execution_id
  544. self.b_objects=[] # objects created by Mantis during execution
  545. class MantisNode:
  546. """
  547. This class contains the basic interface for a Mantis Node.
  548. A MantisNode is used internally by Mantis to represent the final evaluated node graph.
  549. It gets generated with data from a MantisUINode when the graph is read.
  550. """
  551. def __init__(self, signature : tuple,
  552. base_tree : bpy.types.NodeTree,
  553. socket_templates : list[MantisSocketTemplate]=[],):
  554. self.base_tree=base_tree
  555. self.signature = signature
  556. self.inputs = MantisNodeSocketCollection(node=self, is_input=True)
  557. self.outputs = MantisNodeSocketCollection(node=self, is_input=False)
  558. self.parameters = {}
  559. self.drivers = {}
  560. self.node_type='UNINITIALIZED'
  561. self.hierarchy_connections, self.connections = [], []
  562. self.hierarchy_dependencies, self.dependencies = [], []
  563. self.prepared = False
  564. self.executed = False
  565. self.socket_templates = socket_templates
  566. self.mContext = None # for now I am gonna set this at runtime
  567. # I know it isn't "beautiful OOP" or whatever, but it is just easier
  568. # code should be simple and do things in the simplest way.
  569. # in the future I can refactor it, but it will require changes to 100+
  570. # classes, instead of adding about 5 lines of code elsewhere.
  571. if self.socket_templates:
  572. self.init_sockets()
  573. def init_sockets(self) -> None:
  574. self.inputs.init_sockets(self.socket_templates)
  575. self.outputs.init_sockets(self.socket_templates)
  576. def init_parameters(self, additional_parameters = {}) -> None:
  577. for socket in self.inputs:
  578. self.parameters[socket.name] = None
  579. for socket in self.outputs:
  580. self.parameters[socket.name] = None
  581. for key, value in additional_parameters.items():
  582. self.parameters[key]=value
  583. def gen_property_socket_map(self) -> dict:
  584. props_sockets = {}
  585. for s_template in self.socket_templates:
  586. if not s_template.blender_property:
  587. continue
  588. if isinstance(s_template.blender_property, str):
  589. props_sockets[s_template.blender_property]=(s_template.name, s_template.default_value)
  590. elif isinstance(s_template.blender_property, tuple):
  591. for index, sub_prop in enumerate(s_template.blender_property):
  592. props_sockets[sub_prop]=( (s_template.name, index),s_template.default_value[index] )
  593. return props_sockets
  594. def set_traverse(self, traversal_pairs = [(str, str)]) -> None:
  595. for (a, b) in traversal_pairs:
  596. self.inputs[a].set_traverse_target(self.outputs[b])
  597. self.outputs[b].set_traverse_target(self.inputs[a])
  598. def flush_links(self) -> None:
  599. for inp in self.inputs.values():
  600. inp.flush_links()
  601. for out in self.outputs.values():
  602. out.flush_links()
  603. def evaluate_input(self, input_name, index=0) -> Any:
  604. from .node_container_common import trace_single_line
  605. if not (self.inputs.get(input_name)): # get the named parameter if there is no input
  606. return self.parameters.get(input_name) # this will return None if the parameter does not exist.
  607. # this trace() should give a key error if there is a problem
  608. # it is NOT handled here because it should NOT happen - so I want the error message.
  609. trace = trace_single_line(self, input_name, index)
  610. prop = trace[0][-1].parameters[trace[1].name] #trace[0] = the list of traced nodes; read its parameters
  611. return prop
  612. def fill_parameters(self, ui_node=None) -> None:
  613. from .utilities import get_node_prototype
  614. from .node_container_common import get_socket_value
  615. if not ui_node:
  616. if ( (self.signature[0] in ["MANTIS_AUTOGENERATED", "SCHEMA_AUTOGENERATED" ]) or
  617. (self.signature[-1] in ["NodeGroupOutput", "NodeGroupInput"]) ): # I think this is harmless
  618. return None
  619. else:
  620. ui_node = get_node_prototype(self.signature, self.base_tree)
  621. if not ui_node:
  622. raise RuntimeError(wrapRed("No node prototype found for... %s" % ( [self.base_tree] + list(self.signature[1:]) ) ) )
  623. for key in self.parameters.keys():
  624. node_socket = ui_node.inputs.get(key)
  625. if self.parameters[key] is not None: # the parameters are usually initialized as None.
  626. continue # will be filled by the node itself
  627. if not node_socket: #maybe the node socket has no name
  628. if ( ( len(ui_node.inputs) == 0) and ( len(ui_node.outputs) == 1) ):
  629. # this is a simple input node.
  630. node_socket = ui_node.outputs[0]
  631. elif key == 'Name': # for Links we just use the Node Label, or if there is no label, the name.
  632. self.parameters[key] = ui_node.label if ui_node.label else ui_node.name
  633. continue
  634. else:
  635. pass
  636. if node_socket:
  637. if node_socket.bl_idname in ['RelationshipSocket', 'xFormSocket']:
  638. continue
  639. elif node_socket.is_linked and (not node_socket.is_output):
  640. pass # we will get the value from the link, because this is a linked input port.
  641. # very importantly, we do not pass linked outputs- fill these because they are probably Input nodes.
  642. elif hasattr(node_socket, "default_value"):
  643. if (value := get_socket_value(node_socket)) is not None:
  644. self.parameters[key] = value
  645. # TODO: try and remove the input if it is not needed (for performance speed)
  646. else:
  647. raise RuntimeError(wrapRed("No value found for " + self.__repr__() + " when filling out node parameters for " + ui_node.name + "::"+node_socket.name))
  648. else:
  649. pass
  650. # I don't think this works! but I like the idea
  651. def call_on_all_ancestors(self, *args, **kwargs):
  652. """Resolve the dependencies of this node with the named method and its arguments.
  653. First, dependencies are discovered by walking backwards through the tree. Once the root
  654. nodes are discovered, the method is called by each node in dependency order.
  655. The first argument MUST be the name of the method as a string.
  656. """
  657. prGreen(self)
  658. if args[0] == 'call_on_all_ancestors': raise RuntimeError("Very funny!")
  659. from .utilities import get_all_dependencies
  660. from collections import deque
  661. # get all dependencies by walking backward through the tree.
  662. all_dependencies = get_all_dependencies(self)
  663. # get just the roots
  664. can_solve = deque(filter(lambda a : len(a.hierarchy_connections) == 0, all_dependencies))
  665. solved = set()
  666. while can_solve:
  667. node = can_solve.pop()
  668. print(node)
  669. method = getattr(node, args[0])
  670. method(*args[0:], **kwargs)
  671. solved.add(node)
  672. can_solve.extendleft(filter(lambda a : a in all_dependencies, node.hierarchy_connections))
  673. # prPurple(can_solve)
  674. if self in solved:
  675. break
  676. # else:
  677. # for dep in all_dependencies:
  678. # if dep not in solved:
  679. # prOrange(dep)
  680. return
  681. # gets targets for constraints and deformers and should handle all cases
  682. def get_target_and_subtarget(self, constraint_or_deformer, input_name = "Target"):
  683. from bpy.types import PoseBone, Object, SplineIKConstraint
  684. subtarget = ''; target = self.evaluate_input(input_name)
  685. if target:
  686. if not hasattr(target, "bGetObject"):
  687. prRed(f"No {input_name} target found for {constraint_or_deformer.name} in {self} because there is no connected node, or node is wrong type")
  688. return
  689. if (isinstance(target.bGetObject(), PoseBone)):
  690. subtarget = target.bGetObject().name
  691. target = target.bGetParentArmature()
  692. elif (isinstance(target.bGetObject(), Object) ):
  693. target = target.bGetObject()
  694. else:
  695. raise RuntimeError("Cannot interpret constraint or deformer target!")
  696. if (isinstance(constraint_or_deformer, SplineIKConstraint)):
  697. if target and target.type not in ["CURVE"]:
  698. raise GraphError(wrapRed("Error: %s requires a Curve input, not %s" %
  699. (self, type(target))))
  700. constraint_or_deformer.target = target# don't get a subtarget
  701. if (input_name == 'Pole Target'):
  702. constraint_or_deformer.pole_target, constraint_or_deformer.pole_subtarget = target, subtarget
  703. else:
  704. if hasattr(constraint_or_deformer, "target"):
  705. constraint_or_deformer.target = target
  706. if hasattr(constraint_or_deformer, "object"):
  707. constraint_or_deformer.object = target
  708. if hasattr(constraint_or_deformer, "subtarget"):
  709. constraint_or_deformer.subtarget = subtarget
  710. def bPrepare(self, bContext=None):
  711. return
  712. def bExecute(self, bContext=None):
  713. return
  714. def bFinalize(self, bContext=None):
  715. return
  716. def __repr__(self):
  717. return self.signature.__repr__()
  718. # do I need this and the link class above?
  719. class DummyLink:
  720. #gonna use this for faking links to keep the interface consistent
  721. def __init__(self, from_socket, to_socket, nc_from=None, nc_to=None, original_from=None, multi_input_sort_id=0):
  722. self.from_socket = from_socket
  723. self.to_socket = to_socket
  724. self.nc_from = nc_from
  725. self.nc_to = nc_to
  726. self.multi_input_sort_id = multi_input_sort_id
  727. # self.from_node = from_socket.node
  728. # self.to_node = to_socket.node
  729. if (original_from):
  730. self.original_from = original_from
  731. else:
  732. self.original_from = self.from_socket
  733. def __repr__(self):
  734. return(self.nc_from.__repr__()+":"+self.from_socket.name + " -> " + self.nc_to.__repr__()+":"+self.to_socket.name)
  735. def detect_hierarchy_link(from_node, from_socket, to_node, to_socket,):
  736. if to_node.node_type in ['DUMMY_SCHEMA', 'SCHEMA']:
  737. return False
  738. if (from_socket in from_name_filter) or (to_socket in to_name_filter):
  739. return False
  740. # if from_node.__class__.__name__ in ["UtilityCombineVector", "UtilityCombineThreeBool"]:
  741. # return False
  742. return True
  743. class NodeLink:
  744. from_node = None
  745. from_socket = None
  746. to_node = None
  747. to_socket = None
  748. def __init__(self, from_node, from_socket, to_node, to_socket, multi_input_sort_id=0):
  749. if from_node.signature == to_node.signature:
  750. raise RuntimeError("Cannot connect a node to itself.")
  751. self.from_node = from_node
  752. self.from_socket = from_socket
  753. self.to_node = to_node
  754. self.to_socket = to_socket
  755. self.from_node.outputs[self.from_socket].links.append(self)
  756. # it is the responsibility of the node that uses these links to sort them correctly based on the sort_id
  757. self.multi_input_sort_id = multi_input_sort_id
  758. self.to_node.inputs[self.to_socket].links.append(self)
  759. self.is_hierarchy = detect_hierarchy_link(from_node, from_socket, to_node, to_socket,)
  760. self.is_alive = True
  761. def __repr__(self):
  762. return self.from_node.outputs[self.from_socket].__repr__() + " --> " + self.to_node.inputs[self.to_socket].__repr__()
  763. # link_string = # if I need to colorize output for debugging.
  764. # if self.is_hierarchy:
  765. # return wrapOrange(link_string)
  766. # else:
  767. # return wrapWhite(link_string)
  768. def die(self):
  769. self.is_alive = False
  770. self.to_node.inputs[self.to_socket].flush_links()
  771. self.from_node.outputs[self.from_socket].flush_links()
  772. def insert_node(self, middle_node, middle_node_in, middle_node_out, re_init_hierarchy = True):
  773. to_node = self.to_node
  774. to_socket = self.to_socket
  775. self.to_node = middle_node
  776. self.to_socket = middle_node_in
  777. middle_node.outputs[middle_node_out].connect(to_node, to_socket)
  778. if re_init_hierarchy:
  779. from .utilities import init_connections, init_dependencies
  780. init_connections(self.from_node)
  781. init_connections(middle_node)
  782. init_dependencies(middle_node)
  783. init_dependencies(to_node)
  784. class NodeSocket:
  785. # @property # this is a read-only property.
  786. # def is_linked(self):
  787. # return bool(self.links)
  788. def __init__(self, is_input = False,
  789. node = None, name = None,
  790. traverse_target = None):
  791. self.can_traverse = False # to/from the other side of the parent node
  792. self.traverse_target = None
  793. self.node = node
  794. self.name = name
  795. self.is_input = is_input
  796. self.links = []
  797. self.is_linked = False
  798. if (traverse_target):
  799. self.can_traverse = True
  800. def connect(self, node, socket, sort_id=0):
  801. if (self.is_input):
  802. to_node = self.node; from_node = node
  803. to_socket = self.name; from_socket = socket
  804. else:
  805. from_node = self.node; to_node = node
  806. from_socket = self.name; to_socket = socket
  807. from_node.outputs[from_socket].is_linked = True
  808. to_node.inputs[to_socket].is_linked = True
  809. for l in from_node.outputs[from_socket].links:
  810. if l.to_node==to_node and l.to_socket==to_socket:
  811. return None
  812. new_link = NodeLink(
  813. from_node,
  814. from_socket,
  815. to_node,
  816. to_socket,
  817. sort_id)
  818. return new_link
  819. def set_traverse_target(self, traverse_target):
  820. self.traverse_target = traverse_target
  821. self.can_traverse = True
  822. def flush_links(self):
  823. """ Removes dead links from this socket."""
  824. self.links = [l for l in self.links if l.is_alive]
  825. self.links.sort(key=lambda a : -a.multi_input_sort_id)
  826. self.is_linked = bool(self.links)
  827. @property
  828. def is_connected(self):
  829. return len(self.links)>0
  830. def __repr__(self):
  831. return self.node.__repr__() + "::" + self.name
  832. class MantisNodeSocketCollection(dict):
  833. def __init__(self, node, is_input=False):
  834. self.is_input = is_input
  835. self.node = node
  836. def init_sockets(self, sockets):
  837. for socket in sockets:
  838. if isinstance(socket, str):
  839. self[socket] = NodeSocket(is_input=self.is_input, name=socket, node=self.node)
  840. elif isinstance(socket, MantisSocketTemplate):
  841. if socket.is_input != self.is_input: continue
  842. self[socket.name] = NodeSocket(is_input=self.is_input, name=socket.name, node=self.node)
  843. else:
  844. raise RuntimeError(f"NodeSocketCollection keys must be str or MantisSocketTemplate, not {type(socket)}")
  845. def __delitem__(self, key):
  846. """Deletes a node socket by name, and all its links."""
  847. socket = self[key]
  848. for l in socket.links:
  849. l.die()
  850. super().__delitem__(key)
  851. def __iter__(self):
  852. """Makes the class iterable"""
  853. return iter(self.values())