ops_nodegroup.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834
  1. import bpy
  2. from bpy.types import Operator
  3. from mathutils import Vector
  4. from .utilities.utilities import (prRed, prGreen, prPurple, prWhite,
  5. prOrange,
  6. wrapRed, wrapGreen, wrapPurple, wrapWhite,
  7. wrapOrange,)
  8. def mantis_tree_poll_op(context):
  9. space = context.space_data
  10. if hasattr(space, "node_tree"):
  11. if (space.node_tree):
  12. return (space.tree_type in ["MantisTree", "SchemaTree"])
  13. return False
  14. def any_tree_poll(context):
  15. space = context.space_data
  16. if hasattr(space, "node_tree"):
  17. return True
  18. return False
  19. def tag_mantis_node_for_update(context, node):
  20. if any_tree_poll(context):
  21. from .base_definitions import get_signature_from_edited_tree
  22. node_tree = context.space_data.path[0].node_tree
  23. mantis_node = node_tree.parsed_tree.get(
  24. get_signature_from_edited_tree(node, context))
  25. if mantis_node is None:
  26. node_tree.hash = ''
  27. else:
  28. mantis_node.reset_execution_recursive() # tag this as a change.
  29. #########################################################################3
  30. class MantisGroupNodes(Operator):
  31. """Create node-group from selected nodes"""
  32. bl_idname = "mantis.group_nodes"
  33. bl_label = "Group Nodes"
  34. bl_options = {'REGISTER', 'UNDO'}
  35. @classmethod
  36. def poll(cls, context):
  37. return mantis_tree_poll_op(context)
  38. def execute(self, context):
  39. base_tree=context.space_data.path[-1].node_tree
  40. try:
  41. for path_item in context.space_data.path:
  42. path_item.node_tree.is_exporting = True
  43. from .utilities.i_o import export_to_json, do_import
  44. from random import random
  45. grp_name = "".join([chr(int(random()*30)+35) for i in range(20)])
  46. trees=[base_tree]
  47. selected_nodes=export_to_json(trees, base_tree=None, write_file=False, only_selected=True)
  48. # this snippet of confusing indirection edits the name of the base tree in the JSON data
  49. selected_nodes[base_tree.name][0]["name"]=grp_name
  50. do_import(selected_nodes, context)
  51. affected_links_in = []
  52. affected_links_out = []
  53. for l in base_tree.links:
  54. if l.from_node.select and not l.to_node.select: affected_links_out.append(l)
  55. if not l.from_node.select and l.to_node.select: affected_links_in.append(l)
  56. delete_me = []
  57. all_nodes_bounding_box=[Vector((float("inf"),float("inf"))), Vector((-float("inf"),-float("inf")))]
  58. for n in base_tree.nodes:
  59. if n.select:
  60. node_loc = (0,0,0)
  61. if bpy.app.version <= (4, 4):
  62. node_loc = n.location
  63. parent = n.parent
  64. while (parent): # accumulate parent offset
  65. node_loc += parent.location
  66. parent = parent.parent
  67. else: # there is a new location_absolute property in 4.4
  68. node_loc = n.location_absolute
  69. if node_loc.x < all_nodes_bounding_box[0].x:
  70. all_nodes_bounding_box[0].x = node_loc.x
  71. if node_loc.y < all_nodes_bounding_box[0].y:
  72. all_nodes_bounding_box[0].y = node_loc.y
  73. #
  74. if node_loc.x > all_nodes_bounding_box[1].x:
  75. all_nodes_bounding_box[1].x = node_loc.x
  76. if node_loc.y > all_nodes_bounding_box[1].y:
  77. all_nodes_bounding_box[1].y = node_loc.y
  78. delete_me.append(n)
  79. grp_node = base_tree.nodes.new('MantisNodeGroup')
  80. grp_node.node_tree = bpy.data.node_groups[grp_name]
  81. bb_center = all_nodes_bounding_box[0].lerp(all_nodes_bounding_box[1],0.5)
  82. for n in grp_node.node_tree.nodes:
  83. n.location -= bb_center
  84. grp_node.location = Vector((all_nodes_bounding_box[0].x+200, all_nodes_bounding_box[0].lerp(all_nodes_bounding_box[1], 0.5).y))
  85. from .base_definitions import node_group_update
  86. grp_node.is_updating=True
  87. try:
  88. node_group_update(grp_node, force=True)
  89. finally:
  90. grp_node.is_updating=False
  91. # for each node in the JSON
  92. for n in selected_nodes[base_tree.name][2].values():
  93. for socket_collection in ['inputs', 'outputs']:
  94. for s in n[socket_collection].values(): # for each socket in the node
  95. if source := s.get("source"):
  96. prGreen (s["name"], source[0], source[1])
  97. base_tree_node=base_tree.nodes.get(source[0])
  98. if s["is_output"]:
  99. for output in base_tree_node.outputs:
  100. if output.identifier == source[1]:
  101. break
  102. else:
  103. raise RuntimeError(wrapRed("Socket not found when grouping"))
  104. base_tree.links.new(input=output, output=grp_node.inputs[s["name"]])
  105. else:
  106. for s_input in base_tree_node.inputs:
  107. if s_input.identifier == source[1]:
  108. break
  109. else:
  110. raise RuntimeError(wrapRed("Socket not found when grouping"))
  111. base_tree.links.new(input=grp_node.outputs[s["name"]], output=s_input)
  112. for n in delete_me: base_tree.nodes.remove(n)
  113. base_tree.nodes.active = grp_node
  114. finally: # MAKE SURE to turn it back to not exporting
  115. for path_item in context.space_data.path:
  116. path_item.node_tree.is_exporting = False
  117. grp_node.node_tree.name = "Group_Node.000"
  118. return {'FINISHED'}
  119. class MantisEditGroup(Operator):
  120. """Edit the group referenced by the active node (or exit the current node-group)"""
  121. bl_idname = "mantis.edit_group"
  122. bl_label = "Edit Group"
  123. bl_options = {'REGISTER', 'UNDO'}
  124. @classmethod
  125. def poll(cls, context):
  126. return (
  127. mantis_tree_poll_op(context)
  128. )
  129. def execute(self, context):
  130. space = context.space_data
  131. path = space.path
  132. node = path[len(path)-1].node_tree.nodes.active
  133. base_tree = path[0].node_tree
  134. live_update_start_state=base_tree.do_live_update
  135. base_tree.do_live_update = False
  136. base_tree.is_executing = True
  137. try:
  138. if hasattr(node, "node_tree"):
  139. if (node.node_tree):
  140. path.append(node.node_tree, node=node)
  141. path[0].node_tree.display_update(context)
  142. return {"FINISHED"}
  143. elif len(path) > 1:
  144. path.pop()
  145. # get the active node in the current path
  146. active = path[len(path)-1].node_tree.nodes.active
  147. from .base_definitions import node_group_update
  148. active.is_updating = True
  149. try:
  150. node_group_update(active, force = True)
  151. finally:
  152. active.is_updating = False
  153. # call update to force the node group to check if its tree has changed
  154. # now we need to loop through the tree and update all node groups of this type.
  155. from .utilities.utilities import get_all_nodes_of_type
  156. for g in get_all_nodes_of_type(base_tree, "MantisNodeGroup") + \
  157. get_all_nodes_of_type(base_tree, "MantisSchemaGroup"):
  158. if g.node_tree == active.node_tree:
  159. g.is_updating = True
  160. active.is_updating = True
  161. try:
  162. node_group_update(g, force = True)
  163. finally:
  164. g.is_updating = False
  165. active.is_updating = False
  166. base_tree.display_update(context)
  167. base_tree.is_executing = True
  168. # base_tree.is_executing = True # because it seems display_update unsets this.
  169. finally:
  170. base_tree.do_live_update = live_update_start_state
  171. base_tree.is_executing = False
  172. # HACK
  173. base_tree.handler_flip = True # HACK
  174. # HACK
  175. # I have no idea why but the operator finishing causes the exeuction handler to fire
  176. # I have no control over this since it happens after the execution returns...
  177. # so I have to do this ridiculous hack with a Boolean flip bit.
  178. return {"FINISHED"}
  179. class MantisNewNodeTree(Operator):
  180. """Add a new Mantis tree."""
  181. bl_idname = "mantis.new_node_tree"
  182. bl_label = "New Node Group"
  183. bl_options = {'REGISTER', 'UNDO'}
  184. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  185. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  186. @classmethod
  187. def poll(cls, context):
  188. return (
  189. mantis_tree_poll_op(context) and \
  190. context.node.bl_idname in ["MantisNodeGroup", "MantisSchemaGroup"]
  191. )
  192. @classmethod
  193. def poll(cls, context):
  194. return True #(hasattr(context, 'active_node') )
  195. def invoke(self, context, event):
  196. self.tree_invoked = context.node.id_data.name
  197. self.node_invoked = context.node.name
  198. return self.execute(context)
  199. def execute(self, context):
  200. node = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  201. if node.bl_idname == "MantisSchemaGroup":
  202. if (node.node_tree):
  203. print('a')
  204. return {"CANCELLED"}
  205. else:
  206. from bpy import data
  207. print('b')
  208. node.node_tree = data.node_groups.new(name='Schema Group', type='SchemaTree')
  209. return {'FINISHED'}
  210. elif node.bl_idname == "MantisNodeGroup":
  211. if (node.node_tree):
  212. print('c')
  213. return {"CANCELLED"}
  214. else:
  215. from bpy import data
  216. print('d')
  217. node.node_tree = data.node_groups.new(name='Mantis Group', type='MantisTree')
  218. return {'FINISHED'}
  219. else:
  220. return {"CANCELLED"}
  221. class InvalidateNodeTree(Operator):
  222. """Invalidates this node tree, forcing it to read all data again."""
  223. bl_idname = "mantis.invalidate_node_tree"
  224. bl_label = "Clear Node Tree Cache"
  225. bl_options = {'REGISTER', 'UNDO'}
  226. @classmethod
  227. def poll(cls, context):
  228. return (mantis_tree_poll_op(context))
  229. def execute(self, context):
  230. print("Clearing Active Node Tree cache")
  231. tree=context.space_data.path[0].node_tree
  232. tree.execution_id=''; tree.hash=''
  233. return {"FINISHED"}
  234. class ExecuteNodeTree(Operator):
  235. """Execute this node tree"""
  236. bl_idname = "mantis.execute_node_tree"
  237. bl_label = "Execute Node Tree"
  238. bl_options = {'REGISTER', 'UNDO'}
  239. @classmethod
  240. def poll(cls, context):
  241. return (mantis_tree_poll_op(context))
  242. def execute(self, context):
  243. from time import time
  244. from .utilities.utilities import wrapGreen
  245. tree=context.space_data.path[0].node_tree
  246. import cProfile
  247. from os import environ
  248. start_time = time()
  249. do_profile=False
  250. if environ.get("DOPROFILE"):
  251. do_profile=True
  252. pass_error = True
  253. if environ.get("DOERROR"):
  254. pass_error=False
  255. if do_profile:
  256. import pstats, io
  257. from pstats import SortKey
  258. sortby = SortKey.TIME
  259. if 'cum' in environ.get("DOPROFILE").lower():
  260. sortby = SortKey.CUMULATIVE
  261. printme1 = ''
  262. with cProfile.Profile() as pr:
  263. tree.update_tree(context, error_popups = pass_error)
  264. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  265. s1 = io.StringIO()
  266. ps1 = pstats.Stats(pr, stream=s1).strip_dirs().sort_stats(sortby)
  267. ps1.print_stats(30) # print the top 20
  268. # ps1.print_callers("is_signature_in_other_signature")
  269. # ps1.print_callees("is_signature_in_other_signature")
  270. printme1 += s1.getvalue()
  271. printme1 += f"Total profiled time: {ps1.total_tt:.6f} seconds"
  272. with cProfile.Profile() as pr:
  273. tree.execute_tree(context, error_popups = pass_error, profile=True)
  274. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  275. s = io.StringIO()
  276. ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
  277. ps.print_stats(30) # print the top 20
  278. # ps.print_callers("is_signature_in_other_signature")
  279. # ps.print_callees("is_signature_in_other_signature")
  280. print (printme1)
  281. print(s.getvalue())
  282. print(f"Total profiled time: {ps.total_tt:.6f} seconds")
  283. else:
  284. update_failed = tree.update_tree(context, error_popups = pass_error)
  285. if update_failed: # failed. Cannot execute.
  286. return {"FINISHED"}
  287. tree.execute_tree(context, error_popups = pass_error)
  288. prGreen("Finished executing tree in %f seconds" % (time() - start_time))
  289. return {"FINISHED"}
  290. class SelectNodesOfType(Operator):
  291. """Selects all nodes of same type as active node."""
  292. bl_idname = "mantis.select_nodes_of_type"
  293. bl_label = "Select Nodes of Same Type as Active"
  294. bl_options = {'REGISTER', 'UNDO'}
  295. @classmethod
  296. def poll(cls, context):
  297. return (any_tree_poll(context))
  298. def execute(self, context):
  299. active_node = context.active_node
  300. tree = active_node.id_data
  301. if not hasattr(active_node, "node_tree"):
  302. for node in tree.nodes:
  303. node.select = (active_node.bl_idname == node.bl_idname)
  304. else:
  305. for node in tree.nodes:
  306. node.select = (active_node.bl_idname == node.bl_idname) and (active_node.node_tree == node.node_tree)
  307. return {"FINISHED"}
  308. def get_parent_tree_interface_enum(operator, context):
  309. ret = []; i = -1
  310. tree = bpy.data.node_groups[operator.tree_invoked]
  311. for sock in tree.interface.items_tree:
  312. if sock.item_type == 'PANEL': continue
  313. if sock.in_out == "OUTPUT": continue
  314. ret.append( (sock.identifier, sock.name, "Socket from Node Group Input", i := i + 1), )
  315. return ret
  316. def get_node_inputs_enum(operator, context):
  317. ret = []; i = -1
  318. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  319. for inp in n.inputs:
  320. ret.append( (inp.identifier, inp.name, "Socket of node to connect to.", i := i + 1), )
  321. return ret
  322. class ConnectNodeToInput(Operator):
  323. """Connects a Node Group Input socket to specified socket of active node and all selected same-type nodes."""
  324. bl_idname = "mantis.connect_nodes_to_input"
  325. bl_label = "Connect Socket to Input for Selected Nodes"
  326. bl_options = {'REGISTER', 'UNDO'}
  327. group_output : bpy.props.EnumProperty(
  328. items=get_parent_tree_interface_enum,
  329. name="Node Group Input Socket",
  330. description="Select which socket from the Node Group Input to connect to this node",)
  331. node_input : bpy.props.EnumProperty(
  332. items=get_node_inputs_enum,
  333. name="Node Input Socket",
  334. description="Select which of this node's sockets to recieve the connection",)
  335. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  336. world_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  337. material_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  338. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  339. @classmethod
  340. def poll(cls, context):
  341. is_tree = (any_tree_poll(context))
  342. is_material = False
  343. if hasattr(context, 'space_data') and context.space_data.type == 'NODE_EDITOR':
  344. is_material = context.space_data.node_tree.bl_idname == 'ShaderNodeTree'
  345. return is_tree and not is_material # doesn't work for Material right now TODO
  346. def invoke(self, context, event):
  347. self.material_invoked=''; self.world_invoked=''
  348. self.tree_invoked = context.active_node.id_data.name
  349. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  350. if context.space_data.shader_type == 'WORLD':
  351. self.world_invoked = context.space_data.id.name
  352. elif context.space_data.shader_type == 'OBJECT':
  353. self.material_invoked = context.space_data.id.name
  354. else:
  355. return {'CANCELLED'} # what is the LINESTYLE menu? TODO
  356. self.node_invoked = context.active_node.name
  357. # we use active_node here ^ because we are comparing the active node to the selection.
  358. wm = context.window_manager
  359. return wm.invoke_props_dialog(self)
  360. def execute(self, context):
  361. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  362. if context.space_data.shader_type == 'WORLD':
  363. t = bpy.data.worlds[self.material_invoked].node_tree
  364. elif context.space_data.shader_type == 'OBJECT':
  365. t = bpy.data.materials[self.material_invoked].node_tree
  366. else:
  367. t = bpy.data.node_groups.get(self.tree_invoked)
  368. if hasattr(t, "is_executing"): # for Mantis trees, but this function should just work anywhere.
  369. t.is_executing = True
  370. n = t.nodes[self.node_invoked]
  371. for node in t.nodes:
  372. if n.bl_idname == node.bl_idname and node.select:
  373. # the bl_idname is the same so they both have node_tree
  374. if hasattr(n, "node_tree") and n.node_tree != node.node_tree: continue
  375. # TODO: maybe I should try and find a nearby input node and reuse it
  376. # doing these identifier lookups again and again is slow, whatever. faster than doing it by hand
  377. for connect_to_me in node.inputs:
  378. if connect_to_me.identifier == self.node_input: break
  379. if connect_to_me.is_linked: connect_to_me = None
  380. if connect_to_me: # only make the node if the socket is there and free
  381. inp = t.nodes.new("NodeGroupInput")
  382. connect_me = None
  383. for s in inp.outputs:
  384. if s.identifier != self.group_output: s.hide = True
  385. else: connect_me = s
  386. inp.location = node.location
  387. inp.location.x-=200
  388. if connect_me is None:
  389. continue # I don't know how this happens.
  390. # TODO: this bit doesn't work in shader trees for some reason, fix it
  391. t.links.new(input=connect_me, output=connect_to_me)
  392. if hasattr(t, "is_executing"):
  393. t.is_executing = False
  394. return {"FINISHED"}
  395. class QueryNodeSockets(Operator):
  396. """Utility Operator for querying the data in a socket"""
  397. bl_idname = "mantis.query_sockets"
  398. bl_label = "Query Node Sockets"
  399. bl_options = {'REGISTER', 'UNDO'}
  400. @classmethod
  401. def poll(cls, context):
  402. return (mantis_tree_poll_op(context))
  403. def execute(self, context):
  404. active_node = context.active_node
  405. tree = active_node.id_data
  406. for node in tree.nodes:
  407. if not node.select: continue
  408. return {"FINISHED"}
  409. class FoceUpdateGroup(Operator):
  410. """Developer Tool to force a group update"""
  411. bl_idname = "mantis.force_update_group"
  412. bl_label = "Force Mantis Group Update"
  413. bl_options = {'REGISTER', 'UNDO'}
  414. @classmethod
  415. def poll(cls, context):
  416. return (mantis_tree_poll_op(context))
  417. def execute(self, context):
  418. active_node = context.active_node
  419. from .base_definitions import node_group_update
  420. try:
  421. active_node.is_updating = True
  422. node_group_update(active_node, force=True)
  423. finally:
  424. active_node.is_updating = False
  425. return {"FINISHED"}
  426. class ForceDisplayUpdate(Operator):
  427. """Utility Operator for querying the data in a socket"""
  428. bl_idname = "mantis.force_display_update"
  429. bl_label = "Force Mantis Display Update"
  430. bl_options = {'REGISTER', 'UNDO'}
  431. @classmethod
  432. def poll(cls, context):
  433. return (mantis_tree_poll_op(context))
  434. def execute(self, context):
  435. base_tree = bpy.context.space_data.path[0].node_tree
  436. base_tree.display_update(context)
  437. return {"FINISHED"}
  438. class CleanUpNodeGraph(bpy.types.Operator):
  439. """Clean Up Node Graph"""
  440. bl_idname = "mantis.nodes_cleanup"
  441. bl_label = "Clean Up Node Graph"
  442. bl_options = {'REGISTER', 'UNDO'}
  443. # num_iterations=bpy.props.IntProperty(default=8)
  444. @classmethod
  445. def poll(cls, context):
  446. return hasattr(context, 'active_node')
  447. def execute(self, context):
  448. base_tree=context.space_data.path[-1].node_tree
  449. from .utilities.utilities import SugiyamaGraph
  450. SugiyamaGraph(base_tree, 12)
  451. return {'FINISHED'}
  452. class MantisMuteNode(Operator):
  453. """Mantis Test Operator"""
  454. bl_idname = "mantis.mute_node"
  455. bl_label = "Mute Node"
  456. bl_options = {'REGISTER', 'UNDO'}
  457. @classmethod
  458. def poll(cls, context):
  459. return (mantis_tree_poll_op(context))
  460. def execute(self, context):
  461. path = context.space_data.path
  462. node = path[len(path)-1].node_tree.nodes.active
  463. node.mute = not node.mute
  464. # There should only be one of these
  465. if (enable := node.inputs.get("Enable")):
  466. # annoyingly, 'mute' and 'enable' are opposites
  467. enable.default_value = not node.mute
  468. # this one is for Deformers
  469. elif (enable := node.inputs.get("Enable in Viewport")):
  470. # annoyingly, 'mute' and 'enable' are opposites
  471. enable.default_value = not node.mute
  472. elif (hide := node.inputs.get("Hide")):
  473. hide.default_value = node.mute
  474. elif (hide := node.inputs.get("Hide in Viewport")):
  475. hide.default_value = node.mute
  476. return {"FINISHED"}
  477. class DriverAddDriverVariableInput(bpy.types.Operator):
  478. """Add a Driver Variable input to the Driver node"""
  479. bl_idname = "mantis.driver_node_add_variable"
  480. bl_label = "Add Driver Variable"
  481. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  482. @classmethod
  483. def poll(cls, context):
  484. return (hasattr(context, 'active_node') )
  485. def execute(self, context): # unicode for 'a'
  486. i = len (context.node.inputs) - 2 + 96
  487. context.node.inputs.new("DriverVariableSocket", chr(i))
  488. return {'FINISHED'}
  489. class DriverRemoveDriverVariableInput(bpy.types.Operator):
  490. """Remove a DriverVariable input from the active Driver node"""
  491. bl_idname = "mantis.driver_node_remove_variable"
  492. bl_label = "Remove Driver Variable"
  493. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  494. @classmethod
  495. def poll(cls, context):
  496. return (hasattr(context, 'active_node') )
  497. def execute(self, context):
  498. n = context.node
  499. n.inputs.remove(n.inputs[-1])
  500. return {'FINISHED'}
  501. class CollectionAddNewOutput(bpy.types.Operator):
  502. """Add a new Collection output to the Driver node"""
  503. bl_idname = "mantis.collection_add_new"
  504. bl_label = "+ Child"
  505. bl_options = {'REGISTER', 'UNDO'}
  506. collection_name : bpy.props.StringProperty(default='Collection')
  507. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  508. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  509. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  510. @classmethod
  511. def poll(cls, context):
  512. return True #(hasattr(context, 'active_node') )
  513. # DUPLICATED CODE HERE, DUMMY
  514. def invoke(self, context, event):
  515. self.tree_invoked = context.node.id_data.name
  516. self.node_invoked = context.node.name
  517. t = context.node.id_data
  518. t.nodes.active = context.node
  519. context.node.select = True
  520. wm = context.window_manager
  521. return wm.invoke_props_dialog(self)
  522. def execute(self, context):
  523. if not self.collection_name:
  524. return {'CANCELLED'}
  525. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  526. # we need to know which socket it is called from...
  527. s = None
  528. for socket in n.outputs:
  529. if socket.identifier == self.socket_invoked:
  530. s=socket; break
  531. parent_path = ''
  532. if s is not None and s.collection_path:
  533. parent_path = s.collection_path + '>'
  534. outer_dict = n.read_declarations_from_json()
  535. current_data = outer_dict
  536. if parent_path:
  537. for name_elem in parent_path.split('>'):
  538. if name_elem == '': continue # HACK around being a bad programmer
  539. current_data = current_data[name_elem]
  540. current_data[self.collection_name] = {}
  541. n.push_declarations_to_json(outer_dict)
  542. n.update_interface()
  543. return {'FINISHED'}
  544. # TODO: should this prune the children, too?
  545. class CollectionRemoveOutput(bpy.types.Operator):
  546. """Remove a Collection output to the Driver node"""
  547. bl_idname = "mantis.collection_remove"
  548. bl_label = "X"
  549. bl_options = {'REGISTER', 'UNDO'}
  550. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  551. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  552. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  553. @classmethod
  554. def poll(cls, context):
  555. return True #(hasattr(context, 'active_node') )
  556. # DUPLICATED CODE HERE, DUMMY
  557. def invoke(self, context, event):
  558. self.tree_invoked = context.node.id_data.name
  559. self.node_invoked = context.node.name
  560. t = context.node.id_data
  561. t.nodes.active = context.node
  562. context.node.select = True
  563. return self.execute(context)
  564. def execute(self, context):
  565. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  566. s = None
  567. for socket in n.outputs:
  568. if socket.identifier == self.socket_invoked:
  569. s=socket; break
  570. if not s:
  571. return {'CANCELLED'}
  572. parent_path = ''
  573. if s is not None and s.collection_path:
  574. parent_path = s.collection_path + '>'
  575. outer_dict = n.read_declarations_from_json()
  576. current_data = outer_dict
  577. print(parent_path)
  578. if parent_path:
  579. for name_elem in parent_path.split('>')[:-2]: # just skip the last one
  580. print(name_elem)
  581. if name_elem == '': continue # HACK around being a bad programmer
  582. current_data = current_data[name_elem]
  583. del current_data[s.name]
  584. n.push_declarations_to_json(outer_dict)
  585. n.update_interface()
  586. return {'FINISHED'}
  587. def get_socket_enum(operator, context):
  588. valid_types = []; i = -1
  589. from .socket_definitions import TellClasses, MantisSocket
  590. for cls in TellClasses():
  591. if cls.is_valid_interface_type:
  592. valid_types.append( (cls.bl_idname, cls.bl_label, "Socket Type", i := i + 1), )
  593. return valid_types
  594. class B4_4_0_Workaround_NodeTree_Interface_Update(Operator):
  595. """Selects all nodes of same type as active node."""
  596. bl_idname = "mantis.node_tree_interface_update_4_4_0_workaround"
  597. bl_label = "Add Socket to Node Tree"
  598. bl_options = {'REGISTER', 'UNDO'}
  599. socket_name : bpy.props.StringProperty()
  600. output : bpy.props.BoolProperty()
  601. socket_type : bpy.props.EnumProperty(
  602. name="Socket Type",
  603. description="Socket Type",
  604. items=get_socket_enum,
  605. default=0,)
  606. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  607. @classmethod
  608. def poll(cls, context):
  609. return (any_tree_poll(context))
  610. def invoke(self, context, event):
  611. self.tree_invoked = context.active_node.id_data.name
  612. # we use active_node here ^ because we are comparing the active node to the selection.
  613. wm = context.window_manager
  614. return wm.invoke_props_dialog(self)
  615. def execute(self, context):
  616. tree = bpy.data.node_groups[self.tree_invoked]
  617. in_out = 'OUTPUT' if self.output else 'INPUT'
  618. tree.interface.new_socket(self.socket_name, in_out=in_out, socket_type=self.socket_type)
  619. # try to prevent the next execution
  620. # because updating the interface triggers a depsgraph update.
  621. # this doesn't actually work though...TODO
  622. if tree.bl_idname == "MantisTree":
  623. tree.prevent_next_exec=True
  624. return {"FINISHED"}
  625. class ConvertBezierCurveToNURBS(Operator):
  626. """Converts all bezier splines of curve to NURBS."""
  627. bl_idname = "mantis.convert_bezcrv_to_nurbs"
  628. bl_label = "Convert Bezier Curve to NURBS"
  629. bl_options = {'REGISTER', 'UNDO'}
  630. @classmethod
  631. def poll(cls, context):
  632. return (context.active_object is not None) and (context.active_object.type=='CURVE')
  633. def execute(self, context):
  634. from .utilities.utilities import nurbs_copy_bez_spline
  635. curve = context.active_object
  636. bez=[]
  637. for spl in curve.data.splines:
  638. if spl.type=='BEZIER':
  639. bez.append(spl)
  640. for bez_spline in bez:
  641. new_spline=nurbs_copy_bez_spline(curve, bez_spline)
  642. curve.data.splines.remove(bez_spline)
  643. return {"FINISHED"}
  644. def get_component_library_items(self, context):
  645. import json
  646. from .preferences import get_bl_addon_object
  647. bl_mantis_addon = get_bl_addon_object()
  648. return_value = [('NONE', 'None', 'None', 'ERROR', 0)]
  649. component_name_and_path={}
  650. if bl_mantis_addon:
  651. widgets_path = bl_mantis_addon.preferences.ComponentsLibraryFolder
  652. import os
  653. for path_root, dirs, files, in os.walk(widgets_path):
  654. for file in files:
  655. if file.endswith('.rig'):
  656. filepath = os.path.join(path_root, file)
  657. with open(filepath, 'r', encoding='utf-8') as f:
  658. json_data = json.load(f)
  659. for tree_name in json_data.keys():
  660. component_name_and_path[tree_name] = os.path.join(path_root, file)
  661. if component_name_and_path.keys():
  662. return_value=[]
  663. for i, (name, path) in enumerate(component_name_and_path.items()):
  664. return_value.append( ((path+"|"+name).upper(), name, path + ': ' + name, 'NODE_TREE', i) )
  665. return return_value
  666. class ImportFromComponentLibrary(Operator):
  667. """Import a Mantis Component from your Component Library."""
  668. bl_idname = "mantis.import_from_component_library"
  669. bl_label = "Import From Component Library"
  670. bl_options = {'REGISTER', 'UNDO'}
  671. default_value : bpy.props.EnumProperty(
  672. items=get_component_library_items,
  673. name="Component",
  674. description="Which Component to Import",
  675. default = 0,
  676. )
  677. @classmethod
  678. def poll(cls, context):
  679. return True # probably not best to do a True for this LOL
  680. def execute(self, context):
  681. return {"FINISHED"}
  682. from os import environ
  683. if environ.get("DOERROR"): # this should only be available to devs
  684. class TestMantisVersioning(Operator):
  685. """Utility Operator for testing Mantis Versioning"""
  686. bl_idname = "mantis.test_versioning"
  687. bl_label = "Test Mantis Versioning"
  688. bl_options = {'REGISTER', 'UNDO'}
  689. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  690. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  691. @classmethod
  692. def poll(cls, context):
  693. return (mantis_tree_poll_op(context))
  694. def invoke(self, context, event):
  695. self.tree_invoked = bpy.context.space_data.path[-1].node_tree.name
  696. self.node_invoked = bpy.context.space_data.path[-1].node_tree.nodes.active.name
  697. return self.execute(context)
  698. def execute(self, context):
  699. tree = bpy.data.node_groups[self.tree_invoked]
  700. node = tree.nodes[self.node_invoked]
  701. from .versioning import up_0_13_0_new_custom_prop_node
  702. print ("testing...")
  703. up_0_13_0_new_custom_prop_node(node=node, tree=tree)
  704. return {"FINISHED"}
  705. # this has to be down here for some reason. what a pain
  706. classes = [
  707. MantisGroupNodes,
  708. MantisEditGroup,
  709. MantisNewNodeTree,
  710. InvalidateNodeTree,
  711. ExecuteNodeTree,
  712. # CreateMetaGroup,
  713. QueryNodeSockets,
  714. FoceUpdateGroup,
  715. ForceDisplayUpdate,
  716. CleanUpNodeGraph,
  717. MantisMuteNode,
  718. SelectNodesOfType,
  719. ConnectNodeToInput,
  720. # Driver
  721. DriverAddDriverVariableInput,
  722. DriverRemoveDriverVariableInput,
  723. # managing collections
  724. CollectionAddNewOutput,
  725. CollectionRemoveOutput,
  726. # rigging utilities
  727. ConvertBezierCurveToNURBS,
  728. # component library
  729. ImportFromComponentLibrary,
  730. ]
  731. if (bpy.app.version >= (4, 4, 0)):
  732. classes.append(B4_4_0_Workaround_NodeTree_Interface_Update)
  733. if environ.get("DOERROR"):
  734. classes.append(TestMantisVersioning)
  735. def TellClasses():
  736. return classes