base_definitions.py 47 KB

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