ops_nodegroup.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793
  1. import bpy
  2. from bpy.types import Operator
  3. from mathutils import Vector
  4. from .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. #########################################################################3
  20. class MantisGroupNodes(Operator):
  21. """Create node-group from selected nodes"""
  22. bl_idname = "mantis.group_nodes"
  23. bl_label = "Group Nodes"
  24. bl_options = {'REGISTER', 'UNDO'}
  25. @classmethod
  26. def poll(cls, context):
  27. return mantis_tree_poll_op(context)
  28. def execute(self, context):
  29. base_tree=context.space_data.path[-1].node_tree
  30. try:
  31. for path_item in context.space_data.path:
  32. path_item.node_tree.is_exporting = True
  33. from .i_o import export_to_json, do_import
  34. from random import random
  35. grp_name = "".join([chr(int(random()*30)+35) for i in range(20)])
  36. trees=[base_tree]
  37. selected_nodes=export_to_json(trees, write_file=False, only_selected=True)
  38. # this snippet of confusing indirection edits the name of the base tree in the JSON data
  39. selected_nodes[base_tree.name][0]["name"]=grp_name
  40. do_import(selected_nodes, context)
  41. affected_links_in = []
  42. affected_links_out = []
  43. for l in base_tree.links:
  44. if l.from_node.select and not l.to_node.select: affected_links_out.append(l)
  45. if not l.from_node.select and l.to_node.select: affected_links_in.append(l)
  46. delete_me = []
  47. all_nodes_bounding_box=[Vector((float("inf"),float("inf"))), Vector((-float("inf"),-float("inf")))]
  48. for n in base_tree.nodes:
  49. if n.select:
  50. node_loc = (0,0,0)
  51. if bpy.app.version <= (4, 4):
  52. node_loc = n.location
  53. parent = n.parent
  54. while (parent): # accumulate parent offset
  55. node_loc += parent.location
  56. parent = parent.parent
  57. else: # there is a new location_absolute property in 4.4
  58. node_loc = n.location_absolute
  59. if node_loc.x < all_nodes_bounding_box[0].x:
  60. all_nodes_bounding_box[0].x = node_loc.x
  61. if node_loc.y < all_nodes_bounding_box[0].y:
  62. all_nodes_bounding_box[0].y = node_loc.y
  63. #
  64. if node_loc.x > all_nodes_bounding_box[1].x:
  65. all_nodes_bounding_box[1].x = node_loc.x
  66. if node_loc.y > all_nodes_bounding_box[1].y:
  67. all_nodes_bounding_box[1].y = node_loc.y
  68. delete_me.append(n)
  69. grp_node = base_tree.nodes.new('MantisNodeGroup')
  70. grp_node.node_tree = bpy.data.node_groups[grp_name]
  71. bb_center = all_nodes_bounding_box[0].lerp(all_nodes_bounding_box[1],0.5)
  72. for n in grp_node.node_tree.nodes:
  73. n.location -= bb_center
  74. 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))
  75. from .base_definitions import node_group_update
  76. node_group_update(grp_node, force=True)
  77. # for each node in the JSON
  78. for n in selected_nodes[base_tree.name][2].values():
  79. for s in n["sockets"].values(): # for each socket in the node
  80. if source := s.get("source"):
  81. prGreen (s["name"], source[0], source[1])
  82. base_tree_node=base_tree.nodes.get(source[0])
  83. if s["is_output"]:
  84. for output in base_tree_node.outputs:
  85. if output.identifier == source[1]:
  86. break
  87. else:
  88. raise RuntimeError(wrapRed("Socket not found when grouping"))
  89. base_tree.links.new(input=output, output=grp_node.inputs[s["name"]])
  90. else:
  91. for s_input in base_tree_node.inputs:
  92. if s_input.identifier == source[1]:
  93. break
  94. else:
  95. raise RuntimeError(wrapRed("Socket not found when grouping"))
  96. base_tree.links.new(input=grp_node.outputs[s["name"]], output=s_input)
  97. for n in delete_me: base_tree.nodes.remove(n)
  98. base_tree.nodes.active = grp_node
  99. finally: # MAKE SURE to turn it back to not exporting
  100. for path_item in context.space_data.path:
  101. path_item.node_tree.is_exporting = False
  102. grp_node.node_tree.name = "Group_Node.000"
  103. return {'FINISHED'}
  104. class MantisEditGroup(Operator):
  105. """Edit the group referenced by the active node (or exit the current node-group)"""
  106. bl_idname = "mantis.edit_group"
  107. bl_label = "Edit Group"
  108. bl_options = {'REGISTER', 'UNDO'}
  109. @classmethod
  110. def poll(cls, context):
  111. return (
  112. mantis_tree_poll_op(context)
  113. )
  114. def execute(self, context):
  115. space = context.space_data
  116. path = space.path
  117. node = path[len(path)-1].node_tree.nodes.active
  118. base_tree = path[0].node_tree
  119. base_tree.do_live_update = False
  120. base_tree.is_executing = True
  121. try:
  122. if hasattr(node, "node_tree"):
  123. if (node.node_tree):
  124. path.append(node.node_tree, node=node)
  125. path[0].node_tree.display_update(context)
  126. return {"FINISHED"}
  127. elif len(path) > 1:
  128. path.pop()
  129. # get the active node in the current path
  130. active = path[len(path)-1].node_tree.nodes.active
  131. from .base_definitions import node_group_update
  132. active.is_updating = True
  133. try:
  134. node_group_update(active, force = True)
  135. finally:
  136. active.is_updating = False
  137. # call update to force the node group to check if its tree has changed
  138. # now we need to loop through the tree and update all node groups of this type.
  139. from .utilities import get_all_nodes_of_type
  140. for g in get_all_nodes_of_type(base_tree, "MantisNodeGroup"):
  141. if g.node_tree == active.node_tree:
  142. g.is_updating = True
  143. try:
  144. node_group_update(g, force = True)
  145. finally:
  146. g.is_updating = False
  147. base_tree.display_update(context)
  148. base_tree.is_executing = True
  149. # base_tree.is_executing = True # because it seems display_update unsets this.
  150. finally:
  151. base_tree.do_live_update = True
  152. base_tree.is_executing = False
  153. # HACK
  154. base_tree.handler_flip = True # HACK
  155. # HACK
  156. # I have no idea why but the operator finishing causes the exeuction handler to fire
  157. # I have no control over this since it happens after the execution returns...
  158. # so I have to do this ridiculous hack with a Boolean flip bit.
  159. return {"FINISHED"}
  160. class ExecuteNodeTree(Operator):
  161. """Execute this node tree"""
  162. bl_idname = "mantis.execute_node_tree"
  163. bl_label = "Execute Node Tree"
  164. bl_options = {'REGISTER', 'UNDO'}
  165. @classmethod
  166. def poll(cls, context):
  167. return (mantis_tree_poll_op(context))
  168. def execute(self, context):
  169. from time import time
  170. from .utilities import wrapGreen
  171. tree=context.space_data.path[0].node_tree
  172. import cProfile
  173. from os import environ
  174. start_time = time()
  175. do_profile=False
  176. print (environ.get("DOPROFILE"))
  177. if environ.get("DOPROFILE"):
  178. do_profile=True
  179. pass_error = True
  180. if environ.get("DOERROR"):
  181. pass_error=False
  182. if do_profile:
  183. import pstats, io
  184. from pstats import SortKey
  185. with cProfile.Profile() as pr:
  186. tree.update_tree(context)
  187. tree.execute_tree(context, error_popups = pass_error)
  188. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  189. s = io.StringIO()
  190. sortby = SortKey.TIME
  191. # sortby = SortKey.CUMULATIVE
  192. ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
  193. ps.print_stats(20) # print the top 20
  194. print(s.getvalue())
  195. else:
  196. tree.update_tree(context)
  197. tree.execute_tree(context, error_popups = pass_error)
  198. prGreen("Finished executing tree in %f seconds" % (time() - start_time))
  199. return {"FINISHED"}
  200. class SelectNodesOfType(Operator):
  201. """Selects all nodes of same type as active node."""
  202. bl_idname = "mantis.select_nodes_of_type"
  203. bl_label = "Select Nodes of Same Type as Active"
  204. bl_options = {'REGISTER', 'UNDO'}
  205. @classmethod
  206. def poll(cls, context):
  207. return (any_tree_poll(context))
  208. def execute(self, context):
  209. active_node = context.active_node
  210. tree = active_node.id_data
  211. if not hasattr(active_node, "node_tree"):
  212. for node in tree.nodes:
  213. node.select = (active_node.bl_idname == node.bl_idname)
  214. else:
  215. for node in tree.nodes:
  216. node.select = (active_node.bl_idname == node.bl_idname) and (active_node.node_tree == node.node_tree)
  217. return {"FINISHED"}
  218. def get_parent_tree_interface_enum(operator, context):
  219. ret = []; i = -1
  220. tree = bpy.data.node_groups[operator.tree_invoked]
  221. for sock in tree.interface.items_tree:
  222. if sock.item_type == 'PANEL': continue
  223. if sock.in_out == "OUTPUT": continue
  224. ret.append( (sock.identifier, sock.name, "Socket from Node Group Input", i := i + 1), )
  225. return ret
  226. def get_node_inputs_enum(operator, context):
  227. ret = []; i = -1
  228. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  229. for inp in n.inputs:
  230. ret.append( (inp.identifier, inp.name, "Socket of node to connect to.", i := i + 1), )
  231. return ret
  232. class ConnectNodeToInput(Operator):
  233. """Connects a Node Group Input socket to specified socket of active node and all selected same-type nodes."""
  234. bl_idname = "mantis.connect_nodes_to_input"
  235. bl_label = "Connect Socket to Input for Selected Nodes"
  236. bl_options = {'REGISTER', 'UNDO'}
  237. group_output : bpy.props.EnumProperty(
  238. items=get_parent_tree_interface_enum,
  239. name="Node Group Input Socket",
  240. description="Select which socket from the Node Group Input to connect to this node",)
  241. node_input : bpy.props.EnumProperty(
  242. items=get_node_inputs_enum,
  243. name="Node Input Socket",
  244. description="Select which of this node's sockets to recieve the connection",)
  245. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  246. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  247. @classmethod
  248. def poll(cls, context):
  249. return (any_tree_poll(context))
  250. def invoke(self, context, event):
  251. self.tree_invoked = context.active_node.id_data.name
  252. self.node_invoked = context.active_node.name
  253. # we use active_node here ^ because we are comparing the active node to the selection.
  254. wm = context.window_manager
  255. return wm.invoke_props_dialog(self)
  256. def execute(self, context):
  257. t = bpy.data.node_groups[self.tree_invoked]
  258. if hasattr(t, "is_executing"): # for Mantis trees, but this function should just work anywhere.
  259. t.is_executing = True
  260. n = t.nodes[self.node_invoked]
  261. for node in t.nodes:
  262. if n.bl_idname == node.bl_idname and node.select:
  263. # the bl_idname is the same so they both have node_tree
  264. if hasattr(n, "node_tree") and n.node_tree != node.node_tree: continue
  265. # TODO: maybe I should try and find a nearby input node and reuse it
  266. # doing these identifier lookups again and again is slow, whatever. faster than doing it by hand
  267. for connect_to_me in node.inputs:
  268. if connect_to_me.identifier == self.node_input: break
  269. if connect_to_me.is_linked: connect_to_me = None
  270. if connect_to_me: # only make the node if the socket is there and free
  271. inp = t.nodes.new("NodeGroupInput")
  272. connect_me = None
  273. for s in inp.outputs:
  274. if s.identifier != self.group_output: s.hide = True
  275. else: connect_me = s
  276. inp.location = node.location
  277. inp.location.x-=200
  278. t.links.new(input=connect_me, output=connect_to_me)
  279. if hasattr(t, "is_executing"):
  280. t.is_executing = False
  281. return {"FINISHED"}
  282. class QueryNodeSockets(Operator):
  283. """Utility Operator for querying the data in a socket"""
  284. bl_idname = "mantis.query_sockets"
  285. bl_label = "Query Node Sockets"
  286. bl_options = {'REGISTER', 'UNDO'}
  287. @classmethod
  288. def poll(cls, context):
  289. return (mantis_tree_poll_op(context))
  290. def execute(self, context):
  291. active_node = context.active_node
  292. tree = active_node.id_data
  293. for node in tree.nodes:
  294. if not node.select: continue
  295. return {"FINISHED"}
  296. class ForceDisplayUpdate(Operator):
  297. """Utility Operator for querying the data in a socket"""
  298. bl_idname = "mantis.force_display_update"
  299. bl_label = "Force Mantis Display Update"
  300. bl_options = {'REGISTER', 'UNDO'}
  301. @classmethod
  302. def poll(cls, context):
  303. return (mantis_tree_poll_op(context))
  304. def execute(self, context):
  305. base_tree = bpy.context.space_data.path[0].node_tree
  306. base_tree.display_update(context)
  307. return {"FINISHED"}
  308. class CleanUpNodeGraph(bpy.types.Operator):
  309. """Clean Up Node Graph"""
  310. bl_idname = "mantis.nodes_cleanup"
  311. bl_label = "Clean Up Node Graph"
  312. bl_options = {'REGISTER', 'UNDO'}
  313. # num_iterations=bpy.props.IntProperty(default=8)
  314. @classmethod
  315. def poll(cls, context):
  316. return hasattr(context, 'active_node')
  317. def execute(self, context):
  318. base_tree=context.space_data.path[-1].node_tree
  319. from .utilities import SugiyamaGraph
  320. SugiyamaGraph(base_tree, 12)
  321. return {'FINISHED'}
  322. class MantisMuteNode(Operator):
  323. """Mantis Test Operator"""
  324. bl_idname = "mantis.mute_node"
  325. bl_label = "Mute Node"
  326. bl_options = {'REGISTER', 'UNDO'}
  327. @classmethod
  328. def poll(cls, context):
  329. return (mantis_tree_poll_op(context))
  330. def execute(self, context):
  331. path = context.space_data.path
  332. node = path[len(path)-1].node_tree.nodes.active
  333. node.mute = not node.mute
  334. # There should only be one of these
  335. if (enable := node.inputs.get("Enable")):
  336. # annoyingly, 'mute' and 'enable' are opposites
  337. enable.default_value = not node.mute
  338. if (hide := node.inputs.get("Hide")):
  339. hide.default_value = node.mute
  340. return {"FINISHED"}
  341. ePropertyType =(
  342. ('BOOL' , "Boolean", "Boolean", 0),
  343. ('INT' , "Integer", "Integer", 1),
  344. ('FLOAT' , "Float" , "Float" , 2),
  345. ('VECTOR', "Vector" , "Vector" , 3),
  346. ('STRING', "String" , "String" , 4),
  347. #('ENUM' , "Enum" , "Enum" , 5),
  348. )
  349. from .base_definitions import xFormNode
  350. class AddCustomProperty(bpy.types.Operator):
  351. """Add Custom Property to xForm Node"""
  352. bl_idname = "mantis.add_custom_property"
  353. bl_label = "Add Custom Property"
  354. bl_options = {'REGISTER', 'UNDO'}
  355. prop_type : bpy.props.EnumProperty(
  356. items=ePropertyType,
  357. name="New Property Type",
  358. description="Type of data for new Property",
  359. default = 'BOOL',)
  360. prop_name : bpy.props.StringProperty(default='Prop')
  361. min:bpy.props.FloatProperty(default = 0)
  362. max:bpy.props.FloatProperty(default = 1)
  363. soft_min:bpy.props.FloatProperty(default = 0)
  364. soft_max:bpy.props.FloatProperty(default = 1)
  365. description:bpy.props.StringProperty(default = "")
  366. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  367. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  368. @classmethod
  369. def poll(cls, context):
  370. return True #( hasattr(context, 'node') )
  371. def invoke(self, context, event):
  372. self.tree_invoked = context.node.id_data.name
  373. self.node_invoked = context.node.name
  374. wm = context.window_manager
  375. return wm.invoke_props_dialog(self)
  376. def execute(self, context):
  377. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  378. # For whatever reason, context.node doesn't exist anymore
  379. # (probably because I use a window to execute)
  380. # so as a sort of dumb workaround I am saving it to a hidden
  381. # property of the operator... it works but Blender complains.
  382. socktype = ''
  383. if not (self.prop_name):
  384. self.report({'ERROR_INVALID_INPUT'}, "Must name the property.")
  385. return {'CANCELLED'}
  386. if self.prop_type == 'BOOL':
  387. socktype = 'ParameterBoolSocket'
  388. if self.prop_type == 'INT':
  389. socktype = 'ParameterIntSocket'
  390. if self.prop_type == 'FLOAT':
  391. socktype = 'ParameterFloatSocket'
  392. if self.prop_type == 'VECTOR':
  393. socktype = 'ParameterVectorSocket'
  394. if self.prop_type == 'STRING':
  395. socktype = 'ParameterStringSocket'
  396. #if self.prop_type == 'ENUM':
  397. # sock_type = 'ParameterStringSocket'
  398. if (s := n.inputs.get(self.prop_name)):
  399. try:
  400. number = int(self.prop_name[-3:])
  401. # see if it has a number
  402. number+=1
  403. self.prop_name = self.prop_name[:-3] + str(number).zfill(3)
  404. except ValueError:
  405. self.prop_name+='.001'
  406. # WRONG # HACK # TODO # BUG #
  407. new_prop = n.inputs.new( socktype, self.prop_name)
  408. if self.prop_type in ['INT','FLOAT']:
  409. new_prop.min = self.min
  410. new_prop.max = self.max
  411. new_prop.soft_min = self.soft_min
  412. new_prop.soft_max = self.soft_max
  413. new_prop.description = self.description
  414. # now do the output
  415. n.outputs.new( socktype, self.prop_name)
  416. return {'FINISHED'}
  417. def main_get_existing_custom_properties(operator, context):
  418. ret = []; i = -1
  419. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  420. for inp in n.inputs:
  421. if 'Parameter' in inp.bl_idname:
  422. ret.append( (inp.identifier, inp.name, "Custom Property to Modify", i := i + 1), )
  423. return ret
  424. class EditCustomProperty(bpy.types.Operator):
  425. """Edit Custom Property"""
  426. bl_idname = "mantis.edit_custom_property"
  427. bl_label = "Edit Custom Property"
  428. bl_options = {'REGISTER', 'UNDO'}
  429. def get_existing_custom_properties(self, context):
  430. return main_get_existing_custom_properties(self, context)
  431. prop_edit : bpy.props.EnumProperty(
  432. items=get_existing_custom_properties,
  433. name="Property to Edit?",
  434. description="Select which property to edit",)
  435. prop_type : bpy.props.EnumProperty(
  436. items=ePropertyType,
  437. name="New Property Type",
  438. description="Type of data for new Property",
  439. default = 'BOOL',)
  440. prop_name : bpy.props.StringProperty(default='Prop')
  441. min:bpy.props.FloatProperty(default = 0)
  442. max:bpy.props.FloatProperty(default = 1)
  443. soft_min:bpy.props.FloatProperty(default = 0)
  444. soft_max:bpy.props.FloatProperty(default = 1)
  445. description:bpy.props.StringProperty(default = "") # TODO: use getters to fill these automatically
  446. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  447. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  448. @classmethod
  449. def poll(cls, context):
  450. return True #( hasattr(context, 'node') )
  451. def invoke(self, context, event):
  452. self.tree_invoked = context.node.id_data.name
  453. self.node_invoked = context.node.name
  454. wm = context.window_manager
  455. return wm.invoke_props_dialog(self)
  456. def execute(self, context):
  457. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  458. prop = n.inputs.get( self.prop_edit )
  459. if prop:
  460. prop.name = self.prop_name
  461. if (s := n.inputs.get(self.prop_edit)):
  462. if self.prop_type in ['INT','FLOAT']:
  463. prop.min = self.min
  464. prop.max = self.max
  465. prop.soft_min = self.soft_min
  466. prop.soft_max = self.soft_max
  467. prop.description = self.description
  468. return {'FINISHED'}
  469. else:
  470. self.report({'ERROR_INVALID_INPUT'}, "Cannot edit a property that does not exist.")
  471. class RemoveCustomProperty(bpy.types.Operator):
  472. """Remove a Custom Property from an xForm Node"""
  473. bl_idname = "mantis.remove_custom_property"
  474. bl_label = "Remove Custom Property"
  475. bl_options = {'REGISTER', 'UNDO'}
  476. def get_existing_custom_properties(self, context):
  477. return main_get_existing_custom_properties(self, context)
  478. prop_remove : bpy.props.EnumProperty(
  479. items=get_existing_custom_properties,
  480. name="Property to remove?",
  481. description="Select which property to remove",)
  482. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  483. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  484. @classmethod
  485. def poll(cls, context):
  486. return True #(hasattr(context, 'active_node') )
  487. def invoke(self, context, event):
  488. self.tree_invoked = context.node.id_data.name
  489. self.node_invoked = context.node.name
  490. t = context.node.id_data
  491. # HACK the props dialog makes this necesary
  492. # because context.node only exists during the event that
  493. # was created by clicking on the node.
  494. t.nodes.active = context.node # HACK
  495. context.node.select = True # HACK
  496. # I need this bc of the callback for the enum property.
  497. # for whatever reason I can't use tree_invoked there
  498. wm = context.window_manager
  499. return wm.invoke_props_dialog(self)
  500. def execute(self, context):
  501. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  502. # For whatever reason, context.node doesn't exist anymore
  503. # (probably because I use a window to execute)
  504. # so as a sort of dumb workaround I am saving it to a hidden
  505. # property of the operator... it works.
  506. for i, inp in enumerate(n.inputs):
  507. if inp.identifier == self.prop_remove:
  508. break
  509. else:
  510. self.report({'ERROR'}, "Input not found")
  511. return {'CANCELLED'}
  512. # it's possible that the output property's identifier isn't the
  513. # exact same... but I don' care. Shouldn't ever happen. TODO
  514. for j, out in enumerate(n.outputs):
  515. if out.identifier == self.prop_remove:
  516. break
  517. else:
  518. self.report({'ERROR'}, "Output not found")
  519. raise RuntimeError("This should not happen!")
  520. n.inputs.remove ( n.inputs [i] )
  521. n.outputs.remove( n.outputs[j] )
  522. return {'FINISHED'}
  523. # SIMPLE node operators...
  524. # May rewrite these in a more generic way later
  525. class FcurveAddKeyframeInput(bpy.types.Operator):
  526. """Add a keyframe input to the fCurve node"""
  527. bl_idname = "mantis.fcurve_node_add_kf"
  528. bl_label = "Add Keyframe"
  529. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  530. @classmethod
  531. def poll(cls, context):
  532. return (hasattr(context, 'active_node') )
  533. def execute(self, context):
  534. num_keys = len( context.node.inputs)
  535. context.node.inputs.new("KeyframeSocket", "Keyframe."+str(num_keys).zfill(3))
  536. return {'FINISHED'}
  537. class FcurveRemoveKeyframeInput(bpy.types.Operator):
  538. """Remove a keyframe input from the fCurve node"""
  539. bl_idname = "mantis.fcurve_node_remove_kf"
  540. bl_label = "Remove Keyframe"
  541. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  542. @classmethod
  543. def poll(cls, context):
  544. return (hasattr(context, 'active_node') )
  545. def execute(self, context):
  546. n = context.node
  547. n.inputs.remove(n.inputs[-1])
  548. return {'FINISHED'}
  549. class DriverAddDriverVariableInput(bpy.types.Operator):
  550. """Add a Driver Variable input to the Driver node"""
  551. bl_idname = "mantis.driver_node_add_variable"
  552. bl_label = "Add Driver Variable"
  553. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  554. @classmethod
  555. def poll(cls, context):
  556. return (hasattr(context, 'active_node') )
  557. def execute(self, context): # unicode for 'a'
  558. i = len (context.node.inputs) - 2 + 96
  559. context.node.inputs.new("DriverVariableSocket", chr(i))
  560. return {'FINISHED'}
  561. class DriverRemoveDriverVariableInput(bpy.types.Operator):
  562. """Remove a DriverVariable input from the active Driver node"""
  563. bl_idname = "mantis.driver_node_remove_variable"
  564. bl_label = "Remove Driver Variable"
  565. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  566. @classmethod
  567. def poll(cls, context):
  568. return (hasattr(context, 'active_node') )
  569. def execute(self, context):
  570. n = context.node
  571. n.inputs.remove(n.inputs[-1])
  572. return {'FINISHED'}
  573. class LinkArmatureAddTargetInput(bpy.types.Operator):
  574. """Add a Driver Variable input to the Driver node"""
  575. bl_idname = "mantis.link_armature_node_add_target"
  576. bl_label = "Add Target"
  577. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  578. @classmethod
  579. def poll(cls, context):
  580. return hasattr(context, 'node')
  581. def execute(self, context): # unicode for 'a'
  582. num_targets = len( list(context.node.inputs)[6:])//2
  583. context.node.inputs.new("xFormSocket", "Target."+str(num_targets).zfill(3))
  584. context.node.inputs.new("FloatFactorSocket", "Weight."+str(num_targets).zfill(3))
  585. return {'FINISHED'}
  586. class LinkArmatureRemoveTargetInput(bpy.types.Operator):
  587. """Remove a DriverVariable input from the active Driver node"""
  588. bl_idname = "mantis.link_armature_node_remove_target"
  589. bl_label = "Remove Target"
  590. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  591. @classmethod
  592. def poll(cls, context):
  593. return hasattr(context, 'node')
  594. def execute(self, context):
  595. n = context.node
  596. n.inputs.remove(n.inputs[-1]); n.inputs.remove(n.inputs[-1])
  597. return {'FINISHED'}
  598. def get_socket_enum(operator, context):
  599. valid_types = []; i = -1
  600. from .socket_definitions import TellClasses, MantisSocket
  601. for cls in TellClasses():
  602. if cls.is_valid_interface_type:
  603. valid_types.append( (cls.bl_idname, cls.bl_label, "Socket Type", i := i + 1), )
  604. return valid_types
  605. class B4_4_0_Workaround_NodeTree_Interface_Update(Operator):
  606. """Selects all nodes of same type as active node."""
  607. bl_idname = "mantis.node_tree_interface_update_4_4_0_workaround"
  608. bl_label = "Add Socket to Node Tree"
  609. bl_options = {'REGISTER', 'UNDO'}
  610. socket_name : bpy.props.StringProperty()
  611. output : bpy.props.BoolProperty()
  612. socket_type : bpy.props.EnumProperty(
  613. name="Socket Type",
  614. description="Socket Type",
  615. items=get_socket_enum,
  616. default=0,)
  617. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  618. @classmethod
  619. def poll(cls, context):
  620. return (any_tree_poll(context))
  621. def invoke(self, context, event):
  622. self.tree_invoked = context.active_node.id_data.name
  623. # we use active_node here ^ because we are comparing the active node to the selection.
  624. wm = context.window_manager
  625. return wm.invoke_props_dialog(self)
  626. def execute(self, context):
  627. tree = bpy.data.node_groups[self.tree_invoked]
  628. in_out = 'OUTPUT' if self.output else 'INPUT'
  629. tree.interface.new_socket(self.socket_name, in_out=in_out, socket_type=self.socket_type)
  630. # try to prevent the next execution
  631. # because updating the interface triggers a depsgraph update.
  632. # this doesn't actually work though...TODO
  633. if tree.bl_idname == "MantisTree":
  634. tree.prevent_next_exec=True
  635. return {"FINISHED"}
  636. # this has to be down here for some reason. what a pain
  637. classes = [
  638. MantisGroupNodes,
  639. MantisEditGroup,
  640. ExecuteNodeTree,
  641. # CreateMetaGroup,
  642. QueryNodeSockets,
  643. ForceDisplayUpdate,
  644. CleanUpNodeGraph,
  645. MantisMuteNode,
  646. SelectNodesOfType,
  647. ConnectNodeToInput,
  648. # xForm
  649. AddCustomProperty,
  650. EditCustomProperty,
  651. RemoveCustomProperty,
  652. # EditFCurveNode,
  653. FcurveAddKeyframeInput,
  654. FcurveRemoveKeyframeInput,
  655. # Driver
  656. DriverAddDriverVariableInput,
  657. DriverRemoveDriverVariableInput,
  658. # Armature Link Node
  659. LinkArmatureAddTargetInput,
  660. LinkArmatureRemoveTargetInput,
  661. ]
  662. if (bpy.app.version >= (4, 4, 0)):
  663. classes.append(B4_4_0_Workaround_NodeTree_Interface_Update)
  664. def TellClasses():
  665. return classes