ops_nodegroup.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829
  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. printme1 = ''
  260. with cProfile.Profile() as pr:
  261. tree.update_tree(context, error_popups = pass_error)
  262. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  263. s1 = io.StringIO()
  264. # sortby = SortKey.CUMULATIVE
  265. ps1 = pstats.Stats(pr, stream=s1).strip_dirs().sort_stats(sortby)
  266. ps1.print_stats(30) # print the top 20
  267. printme1 += s1.getvalue()
  268. printme1 += f"Total profiled time: {ps1.total_tt:.6f} seconds"
  269. with cProfile.Profile() as pr:
  270. tree.execute_tree(context, error_popups = pass_error, profile=True)
  271. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  272. s = io.StringIO()
  273. # sortby = SortKey.CUMULATIVE
  274. ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
  275. ps.print_stats(30) # print the top 20
  276. print (printme1)
  277. print(s.getvalue())
  278. print(f"Total profiled time: {ps.total_tt:.6f} seconds")
  279. else:
  280. update_failed = tree.update_tree(context, error_popups = pass_error)
  281. if update_failed: # failed. Cannot execute.
  282. return {"FINISHED"}
  283. tree.execute_tree(context, error_popups = pass_error)
  284. prGreen("Finished executing tree in %f seconds" % (time() - start_time))
  285. return {"FINISHED"}
  286. class SelectNodesOfType(Operator):
  287. """Selects all nodes of same type as active node."""
  288. bl_idname = "mantis.select_nodes_of_type"
  289. bl_label = "Select Nodes of Same Type as Active"
  290. bl_options = {'REGISTER', 'UNDO'}
  291. @classmethod
  292. def poll(cls, context):
  293. return (any_tree_poll(context))
  294. def execute(self, context):
  295. active_node = context.active_node
  296. tree = active_node.id_data
  297. if not hasattr(active_node, "node_tree"):
  298. for node in tree.nodes:
  299. node.select = (active_node.bl_idname == node.bl_idname)
  300. else:
  301. for node in tree.nodes:
  302. node.select = (active_node.bl_idname == node.bl_idname) and (active_node.node_tree == node.node_tree)
  303. return {"FINISHED"}
  304. def get_parent_tree_interface_enum(operator, context):
  305. ret = []; i = -1
  306. tree = bpy.data.node_groups[operator.tree_invoked]
  307. for sock in tree.interface.items_tree:
  308. if sock.item_type == 'PANEL': continue
  309. if sock.in_out == "OUTPUT": continue
  310. ret.append( (sock.identifier, sock.name, "Socket from Node Group Input", i := i + 1), )
  311. return ret
  312. def get_node_inputs_enum(operator, context):
  313. ret = []; i = -1
  314. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  315. for inp in n.inputs:
  316. ret.append( (inp.identifier, inp.name, "Socket of node to connect to.", i := i + 1), )
  317. return ret
  318. class ConnectNodeToInput(Operator):
  319. """Connects a Node Group Input socket to specified socket of active node and all selected same-type nodes."""
  320. bl_idname = "mantis.connect_nodes_to_input"
  321. bl_label = "Connect Socket to Input for Selected Nodes"
  322. bl_options = {'REGISTER', 'UNDO'}
  323. group_output : bpy.props.EnumProperty(
  324. items=get_parent_tree_interface_enum,
  325. name="Node Group Input Socket",
  326. description="Select which socket from the Node Group Input to connect to this node",)
  327. node_input : bpy.props.EnumProperty(
  328. items=get_node_inputs_enum,
  329. name="Node Input Socket",
  330. description="Select which of this node's sockets to recieve the connection",)
  331. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  332. world_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  333. material_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  334. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  335. @classmethod
  336. def poll(cls, context):
  337. is_tree = (any_tree_poll(context))
  338. is_material = False
  339. if hasattr(context, 'space_data') and context.space_data.type == 'NODE_EDITOR':
  340. is_material = context.space_data.node_tree.bl_idname == 'ShaderNodeTree'
  341. return is_tree and not is_material # doesn't work for Material right now TODO
  342. def invoke(self, context, event):
  343. self.material_invoked=''; self.world_invoked=''
  344. self.tree_invoked = context.active_node.id_data.name
  345. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  346. if context.space_data.shader_type == 'WORLD':
  347. self.world_invoked = context.space_data.id.name
  348. elif context.space_data.shader_type == 'OBJECT':
  349. self.material_invoked = context.space_data.id.name
  350. else:
  351. return {'CANCELLED'} # what is the LINESTYLE menu? TODO
  352. self.node_invoked = context.active_node.name
  353. # we use active_node here ^ because we are comparing the active node to the selection.
  354. wm = context.window_manager
  355. return wm.invoke_props_dialog(self)
  356. def execute(self, context):
  357. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  358. if context.space_data.shader_type == 'WORLD':
  359. t = bpy.data.worlds[self.material_invoked].node_tree
  360. elif context.space_data.shader_type == 'OBJECT':
  361. t = bpy.data.materials[self.material_invoked].node_tree
  362. else:
  363. t = bpy.data.node_groups.get(self.tree_invoked)
  364. if hasattr(t, "is_executing"): # for Mantis trees, but this function should just work anywhere.
  365. t.is_executing = True
  366. n = t.nodes[self.node_invoked]
  367. for node in t.nodes:
  368. if n.bl_idname == node.bl_idname and node.select:
  369. # the bl_idname is the same so they both have node_tree
  370. if hasattr(n, "node_tree") and n.node_tree != node.node_tree: continue
  371. # TODO: maybe I should try and find a nearby input node and reuse it
  372. # doing these identifier lookups again and again is slow, whatever. faster than doing it by hand
  373. for connect_to_me in node.inputs:
  374. if connect_to_me.identifier == self.node_input: break
  375. if connect_to_me.is_linked: connect_to_me = None
  376. if connect_to_me: # only make the node if the socket is there and free
  377. inp = t.nodes.new("NodeGroupInput")
  378. connect_me = None
  379. for s in inp.outputs:
  380. if s.identifier != self.group_output: s.hide = True
  381. else: connect_me = s
  382. inp.location = node.location
  383. inp.location.x-=200
  384. if connect_me is None:
  385. continue # I don't know how this happens.
  386. # TODO: this bit doesn't work in shader trees for some reason, fix it
  387. t.links.new(input=connect_me, output=connect_to_me)
  388. if hasattr(t, "is_executing"):
  389. t.is_executing = False
  390. return {"FINISHED"}
  391. class QueryNodeSockets(Operator):
  392. """Utility Operator for querying the data in a socket"""
  393. bl_idname = "mantis.query_sockets"
  394. bl_label = "Query Node Sockets"
  395. bl_options = {'REGISTER', 'UNDO'}
  396. @classmethod
  397. def poll(cls, context):
  398. return (mantis_tree_poll_op(context))
  399. def execute(self, context):
  400. active_node = context.active_node
  401. tree = active_node.id_data
  402. for node in tree.nodes:
  403. if not node.select: continue
  404. return {"FINISHED"}
  405. class FoceUpdateGroup(Operator):
  406. """Developer Tool to force a group update"""
  407. bl_idname = "mantis.force_update_group"
  408. bl_label = "Force Mantis Group Update"
  409. bl_options = {'REGISTER', 'UNDO'}
  410. @classmethod
  411. def poll(cls, context):
  412. return (mantis_tree_poll_op(context))
  413. def execute(self, context):
  414. active_node = context.active_node
  415. from .base_definitions import node_group_update
  416. try:
  417. active_node.is_updating = True
  418. node_group_update(active_node, force=True)
  419. finally:
  420. active_node.is_updating = False
  421. return {"FINISHED"}
  422. class ForceDisplayUpdate(Operator):
  423. """Utility Operator for querying the data in a socket"""
  424. bl_idname = "mantis.force_display_update"
  425. bl_label = "Force Mantis Display Update"
  426. bl_options = {'REGISTER', 'UNDO'}
  427. @classmethod
  428. def poll(cls, context):
  429. return (mantis_tree_poll_op(context))
  430. def execute(self, context):
  431. base_tree = bpy.context.space_data.path[0].node_tree
  432. base_tree.display_update(context)
  433. return {"FINISHED"}
  434. class CleanUpNodeGraph(bpy.types.Operator):
  435. """Clean Up Node Graph"""
  436. bl_idname = "mantis.nodes_cleanup"
  437. bl_label = "Clean Up Node Graph"
  438. bl_options = {'REGISTER', 'UNDO'}
  439. # num_iterations=bpy.props.IntProperty(default=8)
  440. @classmethod
  441. def poll(cls, context):
  442. return hasattr(context, 'active_node')
  443. def execute(self, context):
  444. base_tree=context.space_data.path[-1].node_tree
  445. from .utilities.utilities import SugiyamaGraph
  446. SugiyamaGraph(base_tree, 12)
  447. return {'FINISHED'}
  448. class MantisMuteNode(Operator):
  449. """Mantis Test Operator"""
  450. bl_idname = "mantis.mute_node"
  451. bl_label = "Mute Node"
  452. bl_options = {'REGISTER', 'UNDO'}
  453. @classmethod
  454. def poll(cls, context):
  455. return (mantis_tree_poll_op(context))
  456. def execute(self, context):
  457. path = context.space_data.path
  458. node = path[len(path)-1].node_tree.nodes.active
  459. node.mute = not node.mute
  460. # There should only be one of these
  461. if (enable := node.inputs.get("Enable")):
  462. # annoyingly, 'mute' and 'enable' are opposites
  463. enable.default_value = not node.mute
  464. # this one is for Deformers
  465. elif (enable := node.inputs.get("Enable in Viewport")):
  466. # annoyingly, 'mute' and 'enable' are opposites
  467. enable.default_value = not node.mute
  468. elif (hide := node.inputs.get("Hide")):
  469. hide.default_value = node.mute
  470. elif (hide := node.inputs.get("Hide in Viewport")):
  471. hide.default_value = node.mute
  472. return {"FINISHED"}
  473. class DriverAddDriverVariableInput(bpy.types.Operator):
  474. """Add a Driver Variable input to the Driver node"""
  475. bl_idname = "mantis.driver_node_add_variable"
  476. bl_label = "Add Driver Variable"
  477. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  478. @classmethod
  479. def poll(cls, context):
  480. return (hasattr(context, 'active_node') )
  481. def execute(self, context): # unicode for 'a'
  482. i = len (context.node.inputs) - 2 + 96
  483. context.node.inputs.new("DriverVariableSocket", chr(i))
  484. return {'FINISHED'}
  485. class DriverRemoveDriverVariableInput(bpy.types.Operator):
  486. """Remove a DriverVariable input from the active Driver node"""
  487. bl_idname = "mantis.driver_node_remove_variable"
  488. bl_label = "Remove Driver Variable"
  489. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  490. @classmethod
  491. def poll(cls, context):
  492. return (hasattr(context, 'active_node') )
  493. def execute(self, context):
  494. n = context.node
  495. n.inputs.remove(n.inputs[-1])
  496. return {'FINISHED'}
  497. class CollectionAddNewOutput(bpy.types.Operator):
  498. """Add a new Collection output to the Driver node"""
  499. bl_idname = "mantis.collection_add_new"
  500. bl_label = "+ Child"
  501. bl_options = {'REGISTER', 'UNDO'}
  502. collection_name : bpy.props.StringProperty(default='Collection')
  503. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  504. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  505. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  506. @classmethod
  507. def poll(cls, context):
  508. return True #(hasattr(context, 'active_node') )
  509. # DUPLICATED CODE HERE, DUMMY
  510. def invoke(self, context, event):
  511. self.tree_invoked = context.node.id_data.name
  512. self.node_invoked = context.node.name
  513. t = context.node.id_data
  514. t.nodes.active = context.node
  515. context.node.select = True
  516. wm = context.window_manager
  517. return wm.invoke_props_dialog(self)
  518. def execute(self, context):
  519. if not self.collection_name:
  520. return {'CANCELLED'}
  521. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  522. # we need to know which socket it is called from...
  523. s = None
  524. for socket in n.outputs:
  525. if socket.identifier == self.socket_invoked:
  526. s=socket; break
  527. parent_path = ''
  528. if s is not None and s.collection_path:
  529. parent_path = s.collection_path + '>'
  530. outer_dict = n.read_declarations_from_json()
  531. current_data = outer_dict
  532. if parent_path:
  533. for name_elem in parent_path.split('>'):
  534. if name_elem == '': continue # HACK around being a bad programmer
  535. current_data = current_data[name_elem]
  536. current_data[self.collection_name] = {}
  537. n.push_declarations_to_json(outer_dict)
  538. n.update_interface()
  539. return {'FINISHED'}
  540. # TODO: should this prune the children, too?
  541. class CollectionRemoveOutput(bpy.types.Operator):
  542. """Remove a Collection output to the Driver node"""
  543. bl_idname = "mantis.collection_remove"
  544. bl_label = "X"
  545. bl_options = {'REGISTER', 'UNDO'}
  546. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  547. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  548. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  549. @classmethod
  550. def poll(cls, context):
  551. return True #(hasattr(context, 'active_node') )
  552. # DUPLICATED CODE HERE, DUMMY
  553. def invoke(self, context, event):
  554. self.tree_invoked = context.node.id_data.name
  555. self.node_invoked = context.node.name
  556. t = context.node.id_data
  557. t.nodes.active = context.node
  558. context.node.select = True
  559. return self.execute(context)
  560. def execute(self, context):
  561. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  562. s = None
  563. for socket in n.outputs:
  564. if socket.identifier == self.socket_invoked:
  565. s=socket; break
  566. if not s:
  567. return {'CANCELLED'}
  568. parent_path = ''
  569. if s is not None and s.collection_path:
  570. parent_path = s.collection_path + '>'
  571. outer_dict = n.read_declarations_from_json()
  572. current_data = outer_dict
  573. print(parent_path)
  574. if parent_path:
  575. for name_elem in parent_path.split('>')[:-2]: # just skip the last one
  576. print(name_elem)
  577. if name_elem == '': continue # HACK around being a bad programmer
  578. current_data = current_data[name_elem]
  579. del current_data[s.name]
  580. n.push_declarations_to_json(outer_dict)
  581. n.update_interface()
  582. return {'FINISHED'}
  583. def get_socket_enum(operator, context):
  584. valid_types = []; i = -1
  585. from .socket_definitions import TellClasses, MantisSocket
  586. for cls in TellClasses():
  587. if cls.is_valid_interface_type:
  588. valid_types.append( (cls.bl_idname, cls.bl_label, "Socket Type", i := i + 1), )
  589. return valid_types
  590. class B4_4_0_Workaround_NodeTree_Interface_Update(Operator):
  591. """Selects all nodes of same type as active node."""
  592. bl_idname = "mantis.node_tree_interface_update_4_4_0_workaround"
  593. bl_label = "Add Socket to Node Tree"
  594. bl_options = {'REGISTER', 'UNDO'}
  595. socket_name : bpy.props.StringProperty()
  596. output : bpy.props.BoolProperty()
  597. socket_type : bpy.props.EnumProperty(
  598. name="Socket Type",
  599. description="Socket Type",
  600. items=get_socket_enum,
  601. default=0,)
  602. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  603. @classmethod
  604. def poll(cls, context):
  605. return (any_tree_poll(context))
  606. def invoke(self, context, event):
  607. self.tree_invoked = context.active_node.id_data.name
  608. # we use active_node here ^ because we are comparing the active node to the selection.
  609. wm = context.window_manager
  610. return wm.invoke_props_dialog(self)
  611. def execute(self, context):
  612. tree = bpy.data.node_groups[self.tree_invoked]
  613. in_out = 'OUTPUT' if self.output else 'INPUT'
  614. tree.interface.new_socket(self.socket_name, in_out=in_out, socket_type=self.socket_type)
  615. # try to prevent the next execution
  616. # because updating the interface triggers a depsgraph update.
  617. # this doesn't actually work though...TODO
  618. if tree.bl_idname == "MantisTree":
  619. tree.prevent_next_exec=True
  620. return {"FINISHED"}
  621. class ConvertBezierCurveToNURBS(Operator):
  622. """Converts all bezier splines of curve to NURBS."""
  623. bl_idname = "mantis.convert_bezcrv_to_nurbs"
  624. bl_label = "Convert Bezier Curve to NURBS"
  625. bl_options = {'REGISTER', 'UNDO'}
  626. @classmethod
  627. def poll(cls, context):
  628. return (context.active_object is not None) and (context.active_object.type=='CURVE')
  629. def execute(self, context):
  630. from .utilities.utilities import nurbs_copy_bez_spline
  631. curve = context.active_object
  632. bez=[]
  633. for spl in curve.data.splines:
  634. if spl.type=='BEZIER':
  635. bez.append(spl)
  636. for bez_spline in bez:
  637. new_spline=nurbs_copy_bez_spline(curve, bez_spline)
  638. curve.data.splines.remove(bez_spline)
  639. return {"FINISHED"}
  640. def get_component_library_items(self, context):
  641. import json
  642. from .preferences import get_bl_addon_object
  643. bl_mantis_addon = get_bl_addon_object()
  644. return_value = [('NONE', 'None', 'None', 'ERROR', 0)]
  645. component_name_and_path={}
  646. if bl_mantis_addon:
  647. widgets_path = bl_mantis_addon.preferences.ComponentsLibraryFolder
  648. import os
  649. for path_root, dirs, files, in os.walk(widgets_path):
  650. for file in files:
  651. if file.endswith('.rig'):
  652. filepath = os.path.join(path_root, file)
  653. with open(filepath, 'r', encoding='utf-8') as f:
  654. json_data = json.load(f)
  655. for tree_name in json_data.keys():
  656. component_name_and_path[tree_name] = os.path.join(path_root, file)
  657. if component_name_and_path.keys():
  658. return_value=[]
  659. for i, (name, path) in enumerate(component_name_and_path.items()):
  660. return_value.append( ((path+"|"+name).upper(), name, path + ': ' + name, 'NODE_TREE', i) )
  661. return return_value
  662. class ImportFromComponentLibrary(Operator):
  663. """Import a Mantis Component from your Component Library."""
  664. bl_idname = "mantis.import_from_component_library"
  665. bl_label = "Import From Component Library"
  666. bl_options = {'REGISTER', 'UNDO'}
  667. default_value : bpy.props.EnumProperty(
  668. items=get_component_library_items,
  669. name="Component",
  670. description="Which Component to Import",
  671. default = 0,
  672. )
  673. @classmethod
  674. def poll(cls, context):
  675. return True # probably not best to do a True for this LOL
  676. def execute(self, context):
  677. return {"FINISHED"}
  678. from os import environ
  679. if environ.get("DOERROR"): # this should only be available to devs
  680. class TestMantisVersioning(Operator):
  681. """Utility Operator for testing Mantis Versioning"""
  682. bl_idname = "mantis.test_versioning"
  683. bl_label = "Test Mantis Versioning"
  684. bl_options = {'REGISTER', 'UNDO'}
  685. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  686. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  687. @classmethod
  688. def poll(cls, context):
  689. return (mantis_tree_poll_op(context))
  690. def invoke(self, context, event):
  691. self.tree_invoked = bpy.context.space_data.path[-1].node_tree.name
  692. self.node_invoked = bpy.context.space_data.path[-1].node_tree.nodes.active.name
  693. return self.execute(context)
  694. def execute(self, context):
  695. tree = bpy.data.node_groups[self.tree_invoked]
  696. node = tree.nodes[self.node_invoked]
  697. from .versioning import up_0_13_0_new_custom_prop_node
  698. print ("testing...")
  699. up_0_13_0_new_custom_prop_node(node=node, tree=tree)
  700. return {"FINISHED"}
  701. # this has to be down here for some reason. what a pain
  702. classes = [
  703. MantisGroupNodes,
  704. MantisEditGroup,
  705. MantisNewNodeTree,
  706. InvalidateNodeTree,
  707. ExecuteNodeTree,
  708. # CreateMetaGroup,
  709. QueryNodeSockets,
  710. FoceUpdateGroup,
  711. ForceDisplayUpdate,
  712. CleanUpNodeGraph,
  713. MantisMuteNode,
  714. SelectNodesOfType,
  715. ConnectNodeToInput,
  716. # Driver
  717. DriverAddDriverVariableInput,
  718. DriverRemoveDriverVariableInput,
  719. # managing collections
  720. CollectionAddNewOutput,
  721. CollectionRemoveOutput,
  722. # rigging utilities
  723. ConvertBezierCurveToNURBS,
  724. # component library
  725. ImportFromComponentLibrary,
  726. ]
  727. if (bpy.app.version >= (4, 4, 0)):
  728. classes.append(B4_4_0_Workaround_NodeTree_Interface_Update)
  729. if environ.get("DOERROR"):
  730. classes.append(TestMantisVersioning)
  731. def TellClasses():
  732. return classes