ops_nodegroup.py 30 KB

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