ops_nodegroup.py 33 KB

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