ops_nodegroup.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062
  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. grp_node.is_updating=True
  77. try:
  78. node_group_update(grp_node, force=True)
  79. finally:
  80. grp_node.is_updating=False
  81. # for each node in the JSON
  82. for n in selected_nodes[base_tree.name][2].values():
  83. for s in n["sockets"].values(): # for each socket in the node
  84. if source := s.get("source"):
  85. prGreen (s["name"], source[0], source[1])
  86. base_tree_node=base_tree.nodes.get(source[0])
  87. if s["is_output"]:
  88. for output in base_tree_node.outputs:
  89. if output.identifier == source[1]:
  90. break
  91. else:
  92. raise RuntimeError(wrapRed("Socket not found when grouping"))
  93. base_tree.links.new(input=output, output=grp_node.inputs[s["name"]])
  94. else:
  95. for s_input in base_tree_node.inputs:
  96. if s_input.identifier == source[1]:
  97. break
  98. else:
  99. raise RuntimeError(wrapRed("Socket not found when grouping"))
  100. base_tree.links.new(input=grp_node.outputs[s["name"]], output=s_input)
  101. for n in delete_me: base_tree.nodes.remove(n)
  102. base_tree.nodes.active = grp_node
  103. finally: # MAKE SURE to turn it back to not exporting
  104. for path_item in context.space_data.path:
  105. path_item.node_tree.is_exporting = False
  106. grp_node.node_tree.name = "Group_Node.000"
  107. return {'FINISHED'}
  108. class MantisEditGroup(Operator):
  109. """Edit the group referenced by the active node (or exit the current node-group)"""
  110. bl_idname = "mantis.edit_group"
  111. bl_label = "Edit Group"
  112. bl_options = {'REGISTER', 'UNDO'}
  113. @classmethod
  114. def poll(cls, context):
  115. return (
  116. mantis_tree_poll_op(context)
  117. )
  118. def execute(self, context):
  119. space = context.space_data
  120. path = space.path
  121. node = path[len(path)-1].node_tree.nodes.active
  122. base_tree = path[0].node_tree
  123. live_update_start_state=base_tree.do_live_update
  124. base_tree.do_live_update = False
  125. base_tree.is_executing = True
  126. try:
  127. if hasattr(node, "node_tree"):
  128. if (node.node_tree):
  129. path.append(node.node_tree, node=node)
  130. path[0].node_tree.display_update(context)
  131. return {"FINISHED"}
  132. elif len(path) > 1:
  133. path.pop()
  134. # get the active node in the current path
  135. active = path[len(path)-1].node_tree.nodes.active
  136. from .base_definitions import node_group_update
  137. active.is_updating = True
  138. try:
  139. node_group_update(active, force = True)
  140. finally:
  141. active.is_updating = False
  142. # call update to force the node group to check if its tree has changed
  143. # now we need to loop through the tree and update all node groups of this type.
  144. from .utilities import get_all_nodes_of_type
  145. for g in get_all_nodes_of_type(base_tree, "MantisNodeGroup") + \
  146. get_all_nodes_of_type(base_tree, "MantisSchemaGroup"):
  147. if g.node_tree == active.node_tree:
  148. g.is_updating = True
  149. active.is_updating = True
  150. try:
  151. node_group_update(g, force = True)
  152. finally:
  153. g.is_updating = False
  154. active.is_updating = False
  155. base_tree.display_update(context)
  156. base_tree.is_executing = True
  157. # base_tree.is_executing = True # because it seems display_update unsets this.
  158. finally:
  159. base_tree.do_live_update = live_update_start_state
  160. base_tree.is_executing = False
  161. # HACK
  162. base_tree.handler_flip = True # HACK
  163. # HACK
  164. # I have no idea why but the operator finishing causes the exeuction handler to fire
  165. # I have no control over this since it happens after the execution returns...
  166. # so I have to do this ridiculous hack with a Boolean flip bit.
  167. return {"FINISHED"}
  168. class MantisNewNodeTree(Operator):
  169. """Add a new Mantis tree."""
  170. bl_idname = "mantis.new_node_tree"
  171. bl_label = "New Node Group"
  172. bl_options = {'REGISTER', 'UNDO'}
  173. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  174. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  175. @classmethod
  176. def poll(cls, context):
  177. return (
  178. mantis_tree_poll_op(context) and \
  179. context.node.bl_idname in ["MantisNodeGroup", "MantisSchemaGroup"]
  180. )
  181. @classmethod
  182. def poll(cls, context):
  183. return True #(hasattr(context, 'active_node') )
  184. def invoke(self, context, event):
  185. self.tree_invoked = context.node.id_data.name
  186. self.node_invoked = context.node.name
  187. return self.execute(context)
  188. def execute(self, context):
  189. node = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  190. if node.bl_idname == "MantisSchemaGroup":
  191. if (node.node_tree):
  192. print('a')
  193. return {"CANCELLED"}
  194. else:
  195. from bpy import data
  196. print('b')
  197. node.node_tree = data.node_groups.new(name='Schema Group', type='SchemaTree')
  198. return {'FINISHED'}
  199. elif node.bl_idname == "MantisNodeGroup":
  200. if (node.node_tree):
  201. print('c')
  202. return {"CANCELLED"}
  203. else:
  204. from bpy import data
  205. print('d')
  206. node.node_tree = data.node_groups.new(name='Mantis Group', type='MantisTree')
  207. return {'FINISHED'}
  208. else:
  209. return {"CANCELLED"}
  210. class InvalidateNodeTree(Operator):
  211. """Invalidates this node tree, forcing it to read all data again."""
  212. bl_idname = "mantis.invalidate_node_tree"
  213. bl_label = "Clear Node Tree Cache"
  214. bl_options = {'REGISTER', 'UNDO'}
  215. @classmethod
  216. def poll(cls, context):
  217. return (mantis_tree_poll_op(context))
  218. def execute(self, context):
  219. print("Clearing Active Node Tree cache")
  220. tree=context.space_data.path[0].node_tree
  221. tree.execution_id=''; tree.hash=''
  222. return {"FINISHED"}
  223. class ExecuteNodeTree(Operator):
  224. """Execute this node tree"""
  225. bl_idname = "mantis.execute_node_tree"
  226. bl_label = "Execute Node Tree"
  227. bl_options = {'REGISTER', 'UNDO'}
  228. @classmethod
  229. def poll(cls, context):
  230. return (mantis_tree_poll_op(context))
  231. def execute(self, context):
  232. from time import time
  233. from .utilities import wrapGreen
  234. tree=context.space_data.path[0].node_tree
  235. import cProfile
  236. from os import environ
  237. start_time = time()
  238. do_profile=False
  239. if environ.get("DOPROFILE"):
  240. do_profile=True
  241. pass_error = True
  242. if environ.get("DOERROR"):
  243. pass_error=False
  244. if do_profile:
  245. import pstats, io
  246. from pstats import SortKey
  247. with cProfile.Profile() as pr:
  248. tree.update_tree(context, error_popups = pass_error)
  249. tree.execute_tree(context, error_popups = pass_error)
  250. # from the Python docs at https://docs.python.org/3/library/profile.html#module-cProfile
  251. s = io.StringIO()
  252. sortby = SortKey.TIME
  253. # sortby = SortKey.CUMULATIVE
  254. ps = pstats.Stats(pr, stream=s).strip_dirs().sort_stats(sortby)
  255. ps.print_stats(20) # print the top 20
  256. print(s.getvalue())
  257. else:
  258. tree.update_tree(context, error_popups = pass_error)
  259. tree.execute_tree(context, error_popups = pass_error)
  260. prGreen("Finished executing tree in %f seconds" % (time() - start_time))
  261. return {"FINISHED"}
  262. class SelectNodesOfType(Operator):
  263. """Selects all nodes of same type as active node."""
  264. bl_idname = "mantis.select_nodes_of_type"
  265. bl_label = "Select Nodes of Same Type as Active"
  266. bl_options = {'REGISTER', 'UNDO'}
  267. @classmethod
  268. def poll(cls, context):
  269. return (any_tree_poll(context))
  270. def execute(self, context):
  271. active_node = context.active_node
  272. tree = active_node.id_data
  273. if not hasattr(active_node, "node_tree"):
  274. for node in tree.nodes:
  275. node.select = (active_node.bl_idname == node.bl_idname)
  276. else:
  277. for node in tree.nodes:
  278. node.select = (active_node.bl_idname == node.bl_idname) and (active_node.node_tree == node.node_tree)
  279. return {"FINISHED"}
  280. def get_parent_tree_interface_enum(operator, context):
  281. ret = []; i = -1
  282. tree = bpy.data.node_groups[operator.tree_invoked]
  283. for sock in tree.interface.items_tree:
  284. if sock.item_type == 'PANEL': continue
  285. if sock.in_out == "OUTPUT": continue
  286. ret.append( (sock.identifier, sock.name, "Socket from Node Group Input", i := i + 1), )
  287. return ret
  288. def get_node_inputs_enum(operator, context):
  289. ret = []; i = -1
  290. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  291. for inp in n.inputs:
  292. ret.append( (inp.identifier, inp.name, "Socket of node to connect to.", i := i + 1), )
  293. return ret
  294. class ConnectNodeToInput(Operator):
  295. """Connects a Node Group Input socket to specified socket of active node and all selected same-type nodes."""
  296. bl_idname = "mantis.connect_nodes_to_input"
  297. bl_label = "Connect Socket to Input for Selected Nodes"
  298. bl_options = {'REGISTER', 'UNDO'}
  299. group_output : bpy.props.EnumProperty(
  300. items=get_parent_tree_interface_enum,
  301. name="Node Group Input Socket",
  302. description="Select which socket from the Node Group Input to connect to this node",)
  303. node_input : bpy.props.EnumProperty(
  304. items=get_node_inputs_enum,
  305. name="Node Input Socket",
  306. description="Select which of this node's sockets to recieve the connection",)
  307. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  308. world_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  309. material_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  310. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  311. @classmethod
  312. def poll(cls, context):
  313. is_tree = (any_tree_poll(context))
  314. is_material = False
  315. if hasattr(context, 'space_data') and context.space_data.type == 'NODE_EDITOR':
  316. is_material = context.space_data.node_tree.bl_idname == 'ShaderNodeTree'
  317. return is_tree and not is_material # doesn't work for Material right now TODO
  318. def invoke(self, context, event):
  319. self.material_invoked=''; self.world_invoked=''
  320. self.tree_invoked = context.active_node.id_data.name
  321. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  322. if context.space_data.shader_type == 'WORLD':
  323. self.world_invoked = context.space_data.id.name
  324. elif context.space_data.shader_type == 'OBJECT':
  325. self.material_invoked = context.space_data.id.name
  326. else:
  327. return {'CANCELLED'} # what is the LINESTYLE menu? TODO
  328. self.node_invoked = context.active_node.name
  329. # we use active_node here ^ because we are comparing the active node to the selection.
  330. wm = context.window_manager
  331. return wm.invoke_props_dialog(self)
  332. def execute(self, context):
  333. if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
  334. if context.space_data.shader_type == 'WORLD':
  335. t = bpy.data.worlds[self.material_invoked].node_tree
  336. elif context.space_data.shader_type == 'OBJECT':
  337. t = bpy.data.materials[self.material_invoked].node_tree
  338. else:
  339. t = bpy.data.node_groups.get(self.tree_invoked)
  340. if hasattr(t, "is_executing"): # for Mantis trees, but this function should just work anywhere.
  341. t.is_executing = True
  342. n = t.nodes[self.node_invoked]
  343. for node in t.nodes:
  344. if n.bl_idname == node.bl_idname and node.select:
  345. # the bl_idname is the same so they both have node_tree
  346. if hasattr(n, "node_tree") and n.node_tree != node.node_tree: continue
  347. # TODO: maybe I should try and find a nearby input node and reuse it
  348. # doing these identifier lookups again and again is slow, whatever. faster than doing it by hand
  349. for connect_to_me in node.inputs:
  350. if connect_to_me.identifier == self.node_input: break
  351. if connect_to_me.is_linked: connect_to_me = None
  352. if connect_to_me: # only make the node if the socket is there and free
  353. inp = t.nodes.new("NodeGroupInput")
  354. connect_me = None
  355. for s in inp.outputs:
  356. if s.identifier != self.group_output: s.hide = True
  357. else: connect_me = s
  358. inp.location = node.location
  359. inp.location.x-=200
  360. if connect_me is None:
  361. continue # I don't know how this happens.
  362. # TODO: this bit doesn't work in shader trees for some reason, fix it
  363. t.links.new(input=connect_me, output=connect_to_me)
  364. if hasattr(t, "is_executing"):
  365. t.is_executing = False
  366. return {"FINISHED"}
  367. class QueryNodeSockets(Operator):
  368. """Utility Operator for querying the data in a socket"""
  369. bl_idname = "mantis.query_sockets"
  370. bl_label = "Query Node Sockets"
  371. bl_options = {'REGISTER', 'UNDO'}
  372. @classmethod
  373. def poll(cls, context):
  374. return (mantis_tree_poll_op(context))
  375. def execute(self, context):
  376. active_node = context.active_node
  377. tree = active_node.id_data
  378. for node in tree.nodes:
  379. if not node.select: continue
  380. return {"FINISHED"}
  381. class ForceDisplayUpdate(Operator):
  382. """Utility Operator for querying the data in a socket"""
  383. bl_idname = "mantis.force_display_update"
  384. bl_label = "Force Mantis Display Update"
  385. bl_options = {'REGISTER', 'UNDO'}
  386. @classmethod
  387. def poll(cls, context):
  388. return (mantis_tree_poll_op(context))
  389. def execute(self, context):
  390. base_tree = bpy.context.space_data.path[0].node_tree
  391. base_tree.display_update(context)
  392. return {"FINISHED"}
  393. class CleanUpNodeGraph(bpy.types.Operator):
  394. """Clean Up Node Graph"""
  395. bl_idname = "mantis.nodes_cleanup"
  396. bl_label = "Clean Up Node Graph"
  397. bl_options = {'REGISTER', 'UNDO'}
  398. # num_iterations=bpy.props.IntProperty(default=8)
  399. @classmethod
  400. def poll(cls, context):
  401. return hasattr(context, 'active_node')
  402. def execute(self, context):
  403. base_tree=context.space_data.path[-1].node_tree
  404. from .utilities import SugiyamaGraph
  405. SugiyamaGraph(base_tree, 12)
  406. return {'FINISHED'}
  407. class MantisMuteNode(Operator):
  408. """Mantis Test Operator"""
  409. bl_idname = "mantis.mute_node"
  410. bl_label = "Mute Node"
  411. bl_options = {'REGISTER', 'UNDO'}
  412. @classmethod
  413. def poll(cls, context):
  414. return (mantis_tree_poll_op(context))
  415. def execute(self, context):
  416. path = context.space_data.path
  417. node = path[len(path)-1].node_tree.nodes.active
  418. node.mute = not node.mute
  419. # There should only be one of these
  420. if (enable := node.inputs.get("Enable")):
  421. # annoyingly, 'mute' and 'enable' are opposites
  422. enable.default_value = not node.mute
  423. # this one is for Deformers
  424. elif (enable := node.inputs.get("Enable in Viewport")):
  425. # annoyingly, 'mute' and 'enable' are opposites
  426. enable.default_value = not node.mute
  427. elif (hide := node.inputs.get("Hide")):
  428. hide.default_value = node.mute
  429. elif (hide := node.inputs.get("Hide in Viewport")):
  430. hide.default_value = node.mute
  431. return {"FINISHED"}
  432. ePropertyType =(
  433. ('BOOL' , "Boolean", "Boolean", 0),
  434. ('INT' , "Integer", "Integer", 1),
  435. ('FLOAT' , "Float" , "Float" , 2),
  436. ('VECTOR', "Vector" , "Vector" , 3),
  437. ('STRING', "String" , "String" , 4),
  438. #('ENUM' , "Enum" , "Enum" , 5),
  439. )
  440. from .base_definitions import xFormNode
  441. class AddCustomProperty(bpy.types.Operator):
  442. """Add Custom Property to xForm Node"""
  443. bl_idname = "mantis.add_custom_property"
  444. bl_label = "Add Custom Property"
  445. bl_options = {'REGISTER', 'UNDO'}
  446. prop_type : bpy.props.EnumProperty(
  447. items=ePropertyType,
  448. name="New Property Type",
  449. description="Type of data for new Property",
  450. default = 'BOOL',)
  451. prop_name : bpy.props.StringProperty(default='Prop')
  452. min:bpy.props.FloatProperty(default = 0)
  453. max:bpy.props.FloatProperty(default = 1)
  454. soft_min:bpy.props.FloatProperty(default = 0)
  455. soft_max:bpy.props.FloatProperty(default = 1)
  456. description:bpy.props.StringProperty(default = "")
  457. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  458. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  459. @classmethod
  460. def poll(cls, context):
  461. return True #( hasattr(context, 'node') )
  462. def invoke(self, context, event):
  463. self.tree_invoked = context.node.id_data.name
  464. self.node_invoked = context.node.name
  465. wm = context.window_manager
  466. return wm.invoke_props_dialog(self)
  467. def execute(self, context):
  468. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  469. # For whatever reason, context.node doesn't exist anymore
  470. # (probably because I use a window to execute)
  471. # so as a sort of dumb workaround I am saving it to a hidden
  472. # property of the operator... it works but Blender complains.
  473. socktype = ''
  474. if not (self.prop_name):
  475. self.report({'ERROR_INVALID_INPUT'}, "Must name the property.")
  476. return {'CANCELLED'}
  477. if self.prop_type == 'BOOL':
  478. socktype = 'ParameterBoolSocket'
  479. if self.prop_type == 'INT':
  480. socktype = 'ParameterIntSocket'
  481. if self.prop_type == 'FLOAT':
  482. socktype = 'ParameterFloatSocket'
  483. if self.prop_type == 'VECTOR':
  484. socktype = 'ParameterVectorSocket'
  485. if self.prop_type == 'STRING':
  486. socktype = 'ParameterStringSocket'
  487. #if self.prop_type == 'ENUM':
  488. # sock_type = 'ParameterStringSocket'
  489. if (s := n.inputs.get(self.prop_name)):
  490. try:
  491. number = int(self.prop_name[-3:])
  492. # see if it has a number
  493. number+=1
  494. self.prop_name = self.prop_name[:-3] + str(number).zfill(3)
  495. except ValueError:
  496. self.prop_name+='.001'
  497. # WRONG # HACK # TODO # BUG #
  498. new_prop = n.inputs.new( socktype, self.prop_name)
  499. if self.prop_type in ['INT','FLOAT']:
  500. new_prop.min = self.min
  501. new_prop.max = self.max
  502. new_prop.soft_min = self.soft_min
  503. new_prop.soft_max = self.soft_max
  504. new_prop.description = self.description
  505. # now do the output
  506. n.outputs.new( socktype, self.prop_name)
  507. return {'FINISHED'}
  508. def main_get_existing_custom_properties(operator, context):
  509. ret = []; i = -1
  510. n = bpy.data.node_groups[operator.tree_invoked].nodes[operator.node_invoked]
  511. for inp in n.inputs:
  512. if 'Parameter' in inp.bl_idname:
  513. ret.append( (inp.identifier, inp.name, "Custom Property to Modify", i := i + 1), )
  514. return ret
  515. class EditCustomProperty(bpy.types.Operator):
  516. """Edit Custom Property"""
  517. bl_idname = "mantis.edit_custom_property"
  518. bl_label = "Edit Custom Property"
  519. bl_options = {'REGISTER', 'UNDO'}
  520. def get_existing_custom_properties(self, context):
  521. return main_get_existing_custom_properties(self, context)
  522. prop_edit : bpy.props.EnumProperty(
  523. items=get_existing_custom_properties,
  524. name="Property to Edit?",
  525. description="Select which property to edit",)
  526. prop_type : bpy.props.EnumProperty(
  527. items=ePropertyType,
  528. name="New Property Type",
  529. description="Type of data for new Property",
  530. default = 'BOOL',)
  531. prop_name : bpy.props.StringProperty(default='Prop')
  532. min:bpy.props.FloatProperty(default = 0)
  533. max:bpy.props.FloatProperty(default = 1)
  534. soft_min:bpy.props.FloatProperty(default = 0)
  535. soft_max:bpy.props.FloatProperty(default = 1)
  536. description:bpy.props.StringProperty(default = "") # TODO: use getters to fill these automatically
  537. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  538. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  539. @classmethod
  540. def poll(cls, context):
  541. return True #( hasattr(context, 'node') )
  542. def invoke(self, context, event):
  543. self.tree_invoked = context.node.id_data.name
  544. self.node_invoked = context.node.name
  545. wm = context.window_manager
  546. return wm.invoke_props_dialog(self)
  547. def execute(self, context):
  548. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  549. prop = n.inputs.get( self.prop_edit )
  550. if prop:
  551. prop.name = self.prop_name
  552. if (s := n.inputs.get(self.prop_edit)):
  553. if self.prop_type in ['INT','FLOAT']:
  554. prop.min = self.min
  555. prop.max = self.max
  556. prop.soft_min = self.soft_min
  557. prop.soft_max = self.soft_max
  558. prop.description = self.description
  559. return {'FINISHED'}
  560. else:
  561. self.report({'ERROR_INVALID_INPUT'}, "Cannot edit a property that does not exist.")
  562. class RemoveCustomProperty(bpy.types.Operator):
  563. """Remove a Custom Property from an xForm Node"""
  564. bl_idname = "mantis.remove_custom_property"
  565. bl_label = "Remove Custom Property"
  566. bl_options = {'REGISTER', 'UNDO'}
  567. def get_existing_custom_properties(self, context):
  568. return main_get_existing_custom_properties(self, context)
  569. prop_remove : bpy.props.EnumProperty(
  570. items=get_existing_custom_properties,
  571. name="Property to remove?",
  572. description="Select which property to remove",)
  573. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  574. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  575. @classmethod
  576. def poll(cls, context):
  577. return True #(hasattr(context, 'active_node') )
  578. def invoke(self, context, event):
  579. self.tree_invoked = context.node.id_data.name
  580. self.node_invoked = context.node.name
  581. t = context.node.id_data
  582. # HACK the props dialog makes this necesary
  583. # because context.node only exists during the event that
  584. # was created by clicking on the node.
  585. t.nodes.active = context.node # HACK
  586. context.node.select = True # HACK
  587. # I need this bc of the callback for the enum property.
  588. # for whatever reason I can't use tree_invoked there
  589. wm = context.window_manager
  590. return wm.invoke_props_dialog(self)
  591. def execute(self, context):
  592. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  593. # For whatever reason, context.node doesn't exist anymore
  594. # (probably because I use a window to execute)
  595. # so as a sort of dumb workaround I am saving it to a hidden
  596. # property of the operator... it works.
  597. for i, inp in enumerate(n.inputs):
  598. if inp.identifier == self.prop_remove:
  599. break
  600. else:
  601. self.report({'ERROR'}, "Input not found")
  602. return {'CANCELLED'}
  603. # it's possible that the output property's identifier isn't the
  604. # exact same... but I don' care. Shouldn't ever happen. TODO
  605. for j, out in enumerate(n.outputs):
  606. if out.identifier == self.prop_remove:
  607. break
  608. else:
  609. self.report({'ERROR'}, "Output not found")
  610. raise RuntimeError("This should not happen!")
  611. n.inputs.remove ( n.inputs [i] )
  612. n.outputs.remove( n.outputs[j] )
  613. return {'FINISHED'}
  614. # SIMPLE node operators...
  615. # May rewrite these in a more generic way later
  616. class FcurveAddKeyframeInput(bpy.types.Operator):
  617. """Add a keyframe input to the fCurve node"""
  618. bl_idname = "mantis.fcurve_node_add_kf"
  619. bl_label = "Add Keyframe"
  620. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  621. @classmethod
  622. def poll(cls, context):
  623. return (hasattr(context, 'active_node') )
  624. def execute(self, context):
  625. num_keys = len( context.node.inputs)-1
  626. context.node.inputs.new("KeyframeSocket", "Keyframe."+str(num_keys).zfill(3))
  627. return {'FINISHED'}
  628. class FcurveRemoveKeyframeInput(bpy.types.Operator):
  629. """Remove a keyframe input from the fCurve node"""
  630. bl_idname = "mantis.fcurve_node_remove_kf"
  631. bl_label = "Remove Keyframe"
  632. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  633. @classmethod
  634. def poll(cls, context):
  635. return (hasattr(context, 'active_node') )
  636. def execute(self, context):
  637. n = context.node
  638. n.inputs.remove(n.inputs[-1])
  639. return {'FINISHED'}
  640. class DriverAddDriverVariableInput(bpy.types.Operator):
  641. """Add a Driver Variable input to the Driver node"""
  642. bl_idname = "mantis.driver_node_add_variable"
  643. bl_label = "Add Driver Variable"
  644. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  645. @classmethod
  646. def poll(cls, context):
  647. return (hasattr(context, 'active_node') )
  648. def execute(self, context): # unicode for 'a'
  649. i = len (context.node.inputs) - 2 + 96
  650. context.node.inputs.new("DriverVariableSocket", chr(i))
  651. return {'FINISHED'}
  652. class DriverRemoveDriverVariableInput(bpy.types.Operator):
  653. """Remove a DriverVariable input from the active Driver node"""
  654. bl_idname = "mantis.driver_node_remove_variable"
  655. bl_label = "Remove Driver Variable"
  656. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  657. @classmethod
  658. def poll(cls, context):
  659. return (hasattr(context, 'active_node') )
  660. def execute(self, context):
  661. n = context.node
  662. n.inputs.remove(n.inputs[-1])
  663. return {'FINISHED'}
  664. class LinkArmatureAddTargetInput(bpy.types.Operator):
  665. """Add a Driver Variable input to the Driver node"""
  666. bl_idname = "mantis.link_armature_node_add_target"
  667. bl_label = "Add Target"
  668. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  669. @classmethod
  670. def poll(cls, context):
  671. return hasattr(context, 'node')
  672. def execute(self, context): # unicode for 'a'
  673. num_targets = len( list(context.node.inputs)[6:])//2
  674. context.node.inputs.new("xFormSocket", "Target."+str(num_targets).zfill(3))
  675. context.node.inputs.new("FloatFactorSocket", "Weight."+str(num_targets).zfill(3))
  676. return {'FINISHED'}
  677. class LinkArmatureRemoveTargetInput(bpy.types.Operator):
  678. """Remove a DriverVariable input from the active Driver node"""
  679. bl_idname = "mantis.link_armature_node_remove_target"
  680. bl_label = "Remove Target"
  681. bl_options = {'INTERNAL', 'REGISTER', 'UNDO'}
  682. @classmethod
  683. def poll(cls, context):
  684. return hasattr(context, 'node')
  685. def execute(self, context):
  686. n = context.node
  687. n.inputs.remove(n.inputs[-1]); n.inputs.remove(n.inputs[-1])
  688. return {'FINISHED'}
  689. class CollectionAddNewOutput(bpy.types.Operator):
  690. """Add a new Collection output to the Driver node"""
  691. bl_idname = "mantis.collection_add_new"
  692. bl_label = "+ Child"
  693. bl_options = {'REGISTER', 'UNDO'}
  694. collection_name : bpy.props.StringProperty(default='Collection')
  695. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  696. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  697. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  698. @classmethod
  699. def poll(cls, context):
  700. return True #(hasattr(context, 'active_node') )
  701. # DUPLICATED CODE HERE, DUMMY
  702. def invoke(self, context, event):
  703. self.tree_invoked = context.node.id_data.name
  704. self.node_invoked = context.node.name
  705. t = context.node.id_data
  706. t.nodes.active = context.node
  707. context.node.select = True
  708. wm = context.window_manager
  709. return wm.invoke_props_dialog(self)
  710. def execute(self, context):
  711. if not self.collection_name:
  712. return {'CANCELLED'}
  713. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  714. # we need to know which socket it is called from...
  715. s = None
  716. for socket in n.outputs:
  717. if socket.identifier == self.socket_invoked:
  718. s=socket; break
  719. parent_path = ''
  720. if s is not None and s.collection_path:
  721. parent_path = s.collection_path + '>'
  722. outer_dict = n.read_declarations_from_json()
  723. current_data = outer_dict
  724. if parent_path:
  725. for name_elem in parent_path.split('>'):
  726. if name_elem == '': continue # HACK around being a bad programmer
  727. current_data = current_data[name_elem]
  728. current_data[self.collection_name] = {}
  729. n.push_declarations_to_json(outer_dict)
  730. n.update_interface()
  731. return {'FINISHED'}
  732. # TODO: should this prune the children, too?
  733. class CollectionRemoveOutput(bpy.types.Operator):
  734. """Remove a Collection output to the Driver node"""
  735. bl_idname = "mantis.collection_remove"
  736. bl_label = "X"
  737. bl_options = {'REGISTER', 'UNDO'}
  738. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  739. node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  740. socket_invoked : bpy.props.StringProperty(options ={'HIDDEN'}) # set by caller
  741. @classmethod
  742. def poll(cls, context):
  743. return True #(hasattr(context, 'active_node') )
  744. # DUPLICATED CODE HERE, DUMMY
  745. def invoke(self, context, event):
  746. self.tree_invoked = context.node.id_data.name
  747. self.node_invoked = context.node.name
  748. t = context.node.id_data
  749. t.nodes.active = context.node
  750. context.node.select = True
  751. return self.execute(context)
  752. def execute(self, context):
  753. n = bpy.data.node_groups[self.tree_invoked].nodes[self.node_invoked]
  754. s = None
  755. for socket in n.outputs:
  756. if socket.identifier == self.socket_invoked:
  757. s=socket; break
  758. if not s:
  759. return {'CANCELLED'}
  760. parent_path = ''
  761. if s is not None and s.collection_path:
  762. parent_path = s.collection_path + '>'
  763. outer_dict = n.read_declarations_from_json()
  764. current_data = outer_dict
  765. print(parent_path)
  766. if parent_path:
  767. for name_elem in parent_path.split('>')[:-2]: # just skip the last one
  768. print(name_elem)
  769. if name_elem == '': continue # HACK around being a bad programmer
  770. current_data = current_data[name_elem]
  771. del current_data[s.name]
  772. n.push_declarations_to_json(outer_dict)
  773. n.update_interface()
  774. return {'FINISHED'}
  775. def get_socket_enum(operator, context):
  776. valid_types = []; i = -1
  777. from .socket_definitions import TellClasses, MantisSocket
  778. for cls in TellClasses():
  779. if cls.is_valid_interface_type:
  780. valid_types.append( (cls.bl_idname, cls.bl_label, "Socket Type", i := i + 1), )
  781. return valid_types
  782. class B4_4_0_Workaround_NodeTree_Interface_Update(Operator):
  783. """Selects all nodes of same type as active node."""
  784. bl_idname = "mantis.node_tree_interface_update_4_4_0_workaround"
  785. bl_label = "Add Socket to Node Tree"
  786. bl_options = {'REGISTER', 'UNDO'}
  787. socket_name : bpy.props.StringProperty()
  788. output : bpy.props.BoolProperty()
  789. socket_type : bpy.props.EnumProperty(
  790. name="Socket Type",
  791. description="Socket Type",
  792. items=get_socket_enum,
  793. default=0,)
  794. tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
  795. @classmethod
  796. def poll(cls, context):
  797. return (any_tree_poll(context))
  798. def invoke(self, context, event):
  799. self.tree_invoked = context.active_node.id_data.name
  800. # we use active_node here ^ because we are comparing the active node to the selection.
  801. wm = context.window_manager
  802. return wm.invoke_props_dialog(self)
  803. def execute(self, context):
  804. tree = bpy.data.node_groups[self.tree_invoked]
  805. in_out = 'OUTPUT' if self.output else 'INPUT'
  806. tree.interface.new_socket(self.socket_name, in_out=in_out, socket_type=self.socket_type)
  807. # try to prevent the next execution
  808. # because updating the interface triggers a depsgraph update.
  809. # this doesn't actually work though...TODO
  810. if tree.bl_idname == "MantisTree":
  811. tree.prevent_next_exec=True
  812. return {"FINISHED"}
  813. class ConvertBezierCurveToNURBS(Operator):
  814. """Converts all bezier splines of curve to NURBS."""
  815. bl_idname = "mantis.convert_bezcrv_to_nurbs"
  816. bl_label = "Convert Bezier Curve to NURBS"
  817. bl_options = {'REGISTER', 'UNDO'}
  818. @classmethod
  819. def poll(cls, context):
  820. return (context.active_object is not None) and (context.active_object.type=='CURVE')
  821. def execute(self, context):
  822. from .utilities import nurbs_copy_bez_spline
  823. curve = context.active_object
  824. bez=[]
  825. for spl in curve.data.splines:
  826. if spl.type=='BEZIER':
  827. bez.append(spl)
  828. for bez_spline in bez:
  829. new_spline=nurbs_copy_bez_spline(curve, bez_spline)
  830. curve.data.splines.remove(bez_spline)
  831. return {"FINISHED"}
  832. def get_component_library_items(self, context):
  833. import json
  834. from .preferences import get_bl_addon_object
  835. bl_mantis_addon = get_bl_addon_object()
  836. return_value = [('NONE', 'None', 'None', 'ERROR', 0)]
  837. component_name_and_path={}
  838. if bl_mantis_addon:
  839. widgets_path = bl_mantis_addon.preferences.ComponentsLibraryFolder
  840. import os
  841. for path_root, dirs, files, in os.walk(widgets_path):
  842. for file in files:
  843. if file.endswith('.rig'):
  844. filepath = os.path.join(path_root, file)
  845. with open(filepath, 'r', encoding='utf-8') as f:
  846. json_data = json.load(f)
  847. for tree_name in json_data.keys():
  848. component_name_and_path[tree_name] = os.path.join(path_root, file)
  849. if component_name_and_path.keys():
  850. return_value=[]
  851. for i, (name, path) in enumerate(component_name_and_path.items()):
  852. return_value.append( ((path+"|"+name).upper(), name, path + ': ' + name, 'NODE_TREE', i) )
  853. return return_value
  854. class ImportFromComponentLibrary(Operator):
  855. """Import a Mantis Component from your Component Library."""
  856. bl_idname = "mantis.import_from_component_library"
  857. bl_label = "Import From Component Library"
  858. bl_options = {'REGISTER', 'UNDO'}
  859. default_value : bpy.props.EnumProperty(
  860. items=get_component_library_items,
  861. name="Component",
  862. description="Which Component to Import",
  863. default = 0,
  864. )
  865. @classmethod
  866. def poll(cls, context):
  867. return True # probably not best to do a True for this LOL
  868. def execute(self, context):
  869. return {"FINISHED"}
  870. # this has to be down here for some reason. what a pain
  871. classes = [
  872. MantisGroupNodes,
  873. MantisEditGroup,
  874. MantisNewNodeTree,
  875. InvalidateNodeTree,
  876. ExecuteNodeTree,
  877. # CreateMetaGroup,
  878. QueryNodeSockets,
  879. ForceDisplayUpdate,
  880. CleanUpNodeGraph,
  881. MantisMuteNode,
  882. SelectNodesOfType,
  883. ConnectNodeToInput,
  884. # xForm
  885. AddCustomProperty,
  886. EditCustomProperty,
  887. RemoveCustomProperty,
  888. # EditFCurveNode,
  889. FcurveAddKeyframeInput,
  890. FcurveRemoveKeyframeInput,
  891. # Driver
  892. DriverAddDriverVariableInput,
  893. DriverRemoveDriverVariableInput,
  894. # Armature Link Node
  895. LinkArmatureAddTargetInput,
  896. LinkArmatureRemoveTargetInput,
  897. # managing collections
  898. CollectionAddNewOutput,
  899. CollectionRemoveOutput,
  900. # rigging utilities
  901. ConvertBezierCurveToNURBS,
  902. # component library
  903. ImportFromComponentLibrary,
  904. ]
  905. if (bpy.app.version >= (4, 4, 0)):
  906. classes.append(B4_4_0_Workaround_NodeTree_Interface_Update)
  907. def TellClasses():
  908. return classes