utilities.py 66 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519
  1. #fool: should be wrColor like prColor... dumb
  2. def wrapRed(skk): return "\033[91m{}\033[00m".format(skk)
  3. def wrapGreen(skk): return "\033[92m{}\033[00m".format(skk)
  4. def wrapPurple(skk): return "\033[95m{}\033[00m".format(skk)
  5. def wrapWhite(skk): return "\033[97m{}\033[00m".format(skk)
  6. def wrapOrange(skk): return "\033[0;33m{}\033[00m".format(skk)
  7. # these should reimplement the print interface..
  8. def prRed(*args): print (*[wrapRed(arg) for arg in args])
  9. def prGreen(*args): print (*[wrapGreen(arg) for arg in args])
  10. def prPurple(*args): print (*[wrapPurple(arg) for arg in args])
  11. def prWhite(*args): print (*[wrapWhite(arg) for arg in args])
  12. def prOrange(*args): print (*[wrapOrange(arg) for arg in args])
  13. # add THIS to the top of a file for easy access:
  14. # from mantis.utilities import (prRed, prGreen, prPurple, prWhite,
  15. # prOrange,
  16. # wrapRed, wrapGreen, wrapPurple, wrapWhite,
  17. # wrapOrange,)
  18. def float_lerp(a : float, b : float, factor : float) -> float:
  19. return (a * (1.0 - factor)) + (b * factor)
  20. def find_reroute_start_socket(reroute, track='BACK'):
  21. # "BACK" traces back through the tree and returns one socket
  22. # "FORWARD" traces forward through the tree and returns a list of sockets
  23. from ..base_definitions import reroute_cache
  24. from .node_common import GraphError
  25. # We assume the cache is good, the tree is responsible for this
  26. tree = reroute.id_data
  27. tree_cache = reroute_cache.get(tree.name)
  28. if track == 'FORWARD':
  29. cache_item = tree_cache.get(reroute.name); sockets = []
  30. if not cache_item: # I think this is OK, they don't always go somewhere
  31. return # And this FORWARD trace is always a background task
  32. for node_name, socket_index in cache_item.destinations:
  33. sockets.append(tree.nodes.get(node_name).inputs[socket_index])
  34. return sockets
  35. if track == 'BACK':
  36. cache_item = tree_cache.get(reroute.name)
  37. if not cache_item or not cache_item.source:
  38. raise GraphError(f"There is a missing link in the tree at reroute node: {reroute.name}")
  39. return tree.nodes.get(cache_item.source[0]).outputs[cache_item.source[1]]
  40. # this creates fake links that have the same interface as Blender's
  41. # so that I can bypass Reroutes
  42. def clear_reroutes(links):
  43. from ..base_definitions import DummyLink
  44. kept_links, rerouted_starts = [], []
  45. rerouted = []
  46. all_links = links.copy()
  47. while(all_links):
  48. link = all_links.pop()
  49. to_cls = link.to_socket.node.bl_idname
  50. from_cls = link.from_socket.node.bl_idname
  51. reroute_classes = ["NodeReroute"]
  52. if (to_cls in reroute_classes and
  53. from_cls in reroute_classes):
  54. rerouted.append(link)
  55. elif (to_cls in reroute_classes and not
  56. from_cls in reroute_classes):
  57. rerouted.append(link)
  58. elif (from_cls in reroute_classes and not
  59. to_cls in reroute_classes):
  60. rerouted_starts.append(link)
  61. else:
  62. kept_links.append(link)
  63. for start in rerouted_starts:
  64. from_socket = find_reroute_start_socket(start.from_node)
  65. new_link = DummyLink(from_socket=from_socket,
  66. to_socket=start.to_socket,
  67. from_mantis_node=None,
  68. to_mantis_node=None,
  69. multi_input_sort_id=start.multi_input_sort_id )
  70. kept_links.append(new_link)
  71. return kept_links
  72. def tree_from_mantis_node(sig, base_tree):
  73. if (sig[0] == 'MANTIS_AUTOGENERATED'):
  74. sig = sig[:-2] # cut off the end part of the signature (because it uses socket.name and socket.identifier)
  75. # this will lead to totally untraceble bugs in the event of a change in how signatures are assigned
  76. tree = base_tree
  77. for i, path_item in enumerate(sig):
  78. if (i == 0) or (i == len(sig) - 1):
  79. continue
  80. tree = tree.nodes.get(path_item).node_tree
  81. return tree
  82. def get_ui_node(sig, base_tree):
  83. return tree_from_mantis_node(sig, base_tree).nodes.get( sig[-1] )
  84. # Function for splitting up
  85. def split_quoted_substrings_to_strings(text):
  86. import re
  87. pattern = r'''(?:"([^"]*)"|'([^']*)'|([^"']+))'''
  88. substrings = [
  89. next(group for group in match.groups() if group is not None)
  90. for match in re.finditer(pattern, text)
  91. ]
  92. return substrings
  93. def set_bone_selection(pose_bone, selection_state):
  94. # in some older blender versions, select was a bone property
  95. if hasattr(pose_bone, 'select'):
  96. pose_bone.select = selection_state
  97. elif hasattr(pose_bone.bone, 'select'):
  98. pose_bone.bone.select = selection_state
  99. ##################################################################################################
  100. # groups and changing sockets -- this is used extensively by Schema.
  101. ##################################################################################################
  102. # this one returns None if there is an error.
  103. def get_socket_maps(node, force=False):
  104. maps = [{}, {}]
  105. node_collection = ["inputs", "outputs"]
  106. links = ["from_socket", "to_socket"]
  107. for collection, map, linked_socket in zip(node_collection, maps, links):
  108. for sock in getattr(node, collection):
  109. if sock.is_linked:
  110. other_sockets = []
  111. # Sort the links first (in case they are mult-input), because Blender doesn't
  112. links = sorted(list(sock.links), key = lambda l : l.multi_input_sort_id)
  113. # HACK here because Blender will crash if the socket values in the NodeReroute
  114. # are mutated. Because this seems to happen in a deffered way, I can't account
  115. # for it except by checking the node later...
  116. # TODO: The fact that I need this hack means I can probably solve this problem
  117. # for all node types in a safer way, since they may also be dynamic somehow
  118. for l in links:
  119. if "from" in linked_socket and l.from_node.bl_idname == "NodeReroute":
  120. other_sockets.append(l.from_node)
  121. elif "to" in linked_socket and l.to_node.bl_idname == "NodeReroute":
  122. other_sockets.append(l.to_node)
  123. else:
  124. other_sockets.append(getattr(l, linked_socket))
  125. from bpy.types import NodeSocket
  126. keep_sockets=[]
  127. for other_socket in other_sockets.copy():
  128. if isinstance(other_socket, NodeSocket) and \
  129. other_socket.bl_idname == 'NodeSocketUndefined':
  130. continue # this one is bad
  131. keep_sockets.append(other_socket)
  132. map[sock.identifier]= keep_sockets
  133. elif hasattr(sock, "default_value"):
  134. # in blender 5.0 we have to use sock.default_value insted of sock.get("default_value") :/
  135. if sock.bl_idname == "EnumCurveSocket" and sock.default_value is None:
  136. # HACK I need to add this special case because during file-load,
  137. # this value is None and should not be altered until it is set once.
  138. continue
  139. if sock.default_value is not None:
  140. val = getattr(sock, "default_value")
  141. elif "Enum" in sock.bl_idname and isinstance(sock.get("default_value"), int):
  142. continue # for string enum properties that have not yet initialized (at startup)
  143. elif (val := sock.default_value) is not None:
  144. pass
  145. elif not force:
  146. continue
  147. map[sock.identifier]=val
  148. else:
  149. from ..socket_definitions import no_default_value
  150. if sock.bl_idname in no_default_value:
  151. map[sock.identifier]=None
  152. else:
  153. raise RuntimeError(f"ERROR: Could not get socket data for socket of type: {sock.bl_idname}")
  154. return maps
  155. def generate_intervening_reroute(socket, other_socket,):
  156. from bpy.app import version as bpy_version
  157. print ("i have to make a reroute node here lol")
  158. node = socket.node
  159. reroute = node.id_data.nodes.new('NodeReroute')
  160. # NOTE that this doesn't really do a good job of placing it. Whatever.
  161. l = node.id_data.links.new(input=other_socket, output=reroute.inputs[0])
  162. l = node.id_data.links.new(input=reroute.outputs[0], output=socket)
  163. reroute.location = node.location.lerp(other_socket.node.location, 0.25)
  164. if (bpy_version >= (4,4,0)): # correct the reroute color if possible
  165. reroute.socket_idname = other_socket.bl_idname
  166. return l
  167. # this function is completely overloaded with different purposes and code paths
  168. # TODO refactor everything that funnels into this function
  169. # make this stuff simpler.
  170. def do_relink(node, socket, map, in_out='INPUT', parent_name = ''):
  171. if not node.__class__.is_registered_node_type(): return
  172. tree = node.id_data; interface_in_out = 'OUTPUT' if in_out == 'INPUT' else 'INPUT'
  173. if hasattr(node, "node_tree"):
  174. tree = node.node_tree
  175. interface_in_out=in_out
  176. from bpy.types import NodeSocket, Node
  177. get_string = '__extend__'
  178. if socket: get_string = socket.identifier
  179. from ..base_definitions import SchemaUINode
  180. if (hasattr(node, "node_tree") or isinstance(node, SchemaUINode)) and get_string not in map.keys():
  181. # this happens when we are creating a new node group and need to update it from nothing.
  182. return
  183. val = map[get_string] # this will throw an error if the socket isn't there. Good!
  184. # IF there are multiple links going from one node-socket to this socket, then we have to relink
  185. # every socket but the first one before it can go in
  186. # Counter for every socket starts at 0 and if it is greater than 0, add a reroute intervening
  187. # then decrement the counter while relinking,
  188. socket_counters = {}
  189. if isinstance(val, list):
  190. for sub_val in val:
  191. # this will only happen once because it assigns socket, so it is safe to do in the for loop.
  192. if socket is None:
  193. socket = sub_val; sockets = [socket]
  194. if sub_val.bl_idname == "NodeReroute":
  195. # we have to trace the reroute node...
  196. if in_out == 'INPUT':
  197. socket = find_reroute_start_socket(sub_val); sockets = [socket]
  198. else:
  199. sockets = find_reroute_start_socket(sub_val, track="FORWARD")
  200. for socket in sockets:
  201. sock_type = socket.interface_type
  202. name = unique_socket_name(node, socket, tree)
  203. if parent_name:
  204. interface_socket = update_interface(tree.interface, name, interface_in_out, sock_type, parent_name)
  205. if in_out =='INPUT':
  206. socket = node.inputs.new(sock_type, name, identifier=interface_socket.identifier)
  207. else:
  208. socket = node.outputs.new(sock_type, name, identifier=interface_socket.identifier)
  209. if parent_name == 'Array': socket.display_shape='SQUARE_DOT'
  210. if parent_name == 'Constant': socket.display_shape='CIRCLE_DOT'
  211. # then move it up and delete the other link.
  212. # this also needs to modify the interface of the node tree.
  213. if isinstance(sub_val, NodeSocket):
  214. counter_id=sub_val.node.name+"::"+sub_val.identifier
  215. l = None
  216. if in_out =='INPUT':
  217. if socket_counters.get(counter_id) is None:
  218. socket_counters[counter_id] = 0
  219. if socket_counters[counter_id] == 0:
  220. l = node.id_data.links.new(input=sub_val, output=socket)
  221. else:
  222. l = generate_intervening_reroute(socket, sub_val)
  223. socket_counters[counter_id] += 1
  224. else:
  225. l = node.id_data.links.new(input=socket, output=sub_val)
  226. if l is None:
  227. raise RuntimeError("Could not create link")
  228. elif isinstance(sub_val, Node):
  229. counter_id=sub_val.name+"::Output"
  230. l = None
  231. # this happens when it is a NodeReroute
  232. if not socket.is_output:
  233. if socket_counters.get(counter_id) is None:
  234. socket_counters[counter_id] = 0
  235. if socket_counters[counter_id] == 0:
  236. l = node.id_data.links.new(input=sub_val.outputs[0], output=socket)
  237. else:
  238. # NOTE maybe this should make another reroute from the source to here instead
  239. l = generate_intervening_reroute(socket, sub_val.outputs[0])
  240. socket_counters[counter_id] += 1
  241. else:
  242. l = node.id_data.links.new(input=socket, output=sub_val.inputs[0])
  243. if l is None:
  244. raise RuntimeError("Could not create link")
  245. else:
  246. raise RuntimeError("Unhandled case in do_relink()")
  247. elif get_string != "__extend__":
  248. if not socket.is_output:
  249. from bpy.app import version as bpy_version
  250. if bpy_version >=(4,5,0): # VERSIONING
  251. # for some reason, this is throwing an error now
  252. from bpy.types import bpy_prop_array
  253. if isinstance(val, bpy_prop_array):
  254. if in_out == "INPUT" and hasattr(socket, 'input') and socket.input == False:
  255. return # doesn't matter, this is a Matrix socket in a bone or something
  256. # raise RuntimeError(
  257. # f"Cannot set property in socket of type {socket.bl_idname} due to bug in Blender: "
  258. # f"{node.id_data.name}:{node.name}:{socket.name} ")
  259. # TODO: report this weird bug!
  260. try:
  261. if socket.bl_idname == 'BooleanThreeTupleSocket':
  262. # it is so annoying that I have to do this
  263. socket.default_value = [bool(val[0]), bool(val[1]), bool(val[2])]
  264. else:
  265. socket.default_value = val
  266. except (AttributeError, ValueError): # must be readonly or maybe it doesn't have a d.v.
  267. pass
  268. def read_schema_type(interface_item):
  269. # VERSIONING CODE
  270. tree=interface_item.id_data
  271. version = tree.mantis_version
  272. old_version = False
  273. if version[0] == 0:
  274. if version[1] < 13: old_version = True
  275. # unfortunately we need to check this stuff for the versioning code to run correctly the first time.
  276. # UNLESS I can find a way to prevent this code from running before versioning
  277. if old_version or (not hasattr(interface_item, 'is_array')):
  278. # it is not a custom interface class and/or the file is old.
  279. if interface_item.parent:
  280. return interface_item.parent.name
  281. else:
  282. if interface_item.is_array:
  283. return 'Array'
  284. if interface_item.is_connection:
  285. return 'Connection'
  286. return 'Constant'
  287. def update_interface(interface, name, in_out, sock_type, parent_name):
  288. from bpy.app import version as bpy_version
  289. if parent_name:
  290. if not (interface_parent := interface.items_tree.get(parent_name)):
  291. interface_parent = interface.new_panel(name=parent_name)
  292. if bpy_version != (4,5,0):
  293. socket = interface.new_socket(name=name,in_out=in_out, socket_type=sock_type, parent=interface_parent)
  294. else: # blender 4.5.0 LTS, have to workaround a bug!
  295. from ..versioning import workaround_4_5_0_interface_update
  296. socket = workaround_4_5_0_interface_update(tree=interface.id_data, name=name, in_out=in_out,
  297. sock_type=sock_type, parent_name=parent_name, do_parent=True)
  298. match parent_name:
  299. case "Connection":
  300. socket.is_connection=True
  301. case "Array":
  302. socket.is_array=True
  303. if parent_name == 'Connection':
  304. in_out = 'OUTPUT' if in_out == 'INPUT' else 'INPUT' # flip this make sure connections always do both
  305. other_socket = interface.new_socket(name=name,in_out=in_out, socket_type=sock_type, parent=interface_parent)
  306. other_socket.is_connection = True
  307. return socket
  308. else:
  309. raise RuntimeError(wrapRed("Cannot add interface item to tree without specifying type."))
  310. #UGLY BAD REFACTOR
  311. def relink_socket_map_add_socket(node, socket_collection, item, in_out=None,):
  312. from bpy.app import version as bpy_version
  313. # if not in_out: in_out=item.in_out
  314. multi=False
  315. if in_out == 'INPUT' and read_schema_type(item) == 'Array':
  316. multi = True
  317. # have to work around a bug in 4.5.0 that prevents me from declaring custom socket types
  318. # I have arbitrarily chosen to use the NodeSocketGeometry type to signal that this one is affected.
  319. if bpy_version == (4, 5, 0) and item.bl_socket_idname == 'NodeSocketGeometry':
  320. from ..versioning import socket_add_workaround_for_4_5_0_LTS
  321. s = socket_add_workaround_for_4_5_0_LTS(item, socket_collection, multi)
  322. else:
  323. s = socket_collection.new(type=item.bl_socket_idname, name=item.name, identifier=item.identifier, use_multi_input=multi)
  324. if hasattr(s, 'default_value') and hasattr(s, 'is_valid_interface_type') and \
  325. s.is_valid_interface_type == True:
  326. if s.bl_idname not in ['MatrixSocket']: # no default value implemented
  327. from bpy.types import bpy_prop_array
  328. from mathutils import Vector
  329. default_value = 'REPORT BUG ON GITLAB' # default to bug string
  330. val_type = type(s.default_value) # why tf can't I match/case here?
  331. if val_type is bool: default_value = item.default_bool
  332. if val_type is int: default_value = item.default_int
  333. if val_type is float: default_value = item.default_float
  334. if val_type is Vector: default_value = item.default_vector
  335. if val_type is str: default_value = item.default_string
  336. if val_type is bpy_prop_array: default_value = item.default_bool_vector
  337. if not s.is_property_readonly("default_value"):
  338. s.default_value = default_value
  339. else: # TODO: make this work with color sets.
  340. raise Warning("WARN: Not yet implemented: default values for color sets.")
  341. if read_schema_type(item) == 'Array': s.display_shape = 'SQUARE_DOT'
  342. elif node.bl_idname in ['MantisSchemaGroup'] and read_schema_type(item) == 'Constant':
  343. s.display_shape='CIRCLE_DOT'
  344. return s
  345. # TODO REFACTOR THIS
  346. # I did this awful thing because I needed the above code
  347. # but I have provided this interface to Mantis
  348. # I did not follow the Single Responsibility Principle
  349. # I am now suffering for it, as I rightly deserve.
  350. def relink_socket_map(node, socket_collection, map, item, in_out):
  351. new_socket = relink_socket_map_add_socket(node, socket_collection, item, in_out,)
  352. do_relink(node, new_socket, map, in_out, parent_name=read_schema_type(item))
  353. def unique_socket_name(node, other_socket, tree):
  354. name_stem = other_socket.bl_label; num=0
  355. # if hasattr(other_socket, "default_value"):
  356. # name_stem = type(other_socket.default_value).__name__
  357. for item in tree.interface.items_tree:
  358. if item.item_type == 'PANEL': continue
  359. if other_socket.is_output and item.in_out == 'INPUT': continue
  360. if not other_socket.is_output and item.in_out == 'OUTPUT': continue
  361. if name_stem in item.name: num+=1
  362. name = name_stem + '.' + str(num).zfill(3)
  363. return name
  364. ##############################
  365. # Dealing with Objects
  366. ##############################
  367. # use this to ensure the active object is set back when changing it
  368. def preserve_active_object(func):
  369. def wrapper(*args, **kwargs):
  370. import bpy
  371. original_active = bpy.context.active_object
  372. func(*args, **kwargs)
  373. bpy.context.view_layer.objects.active = original_active
  374. return wrapper
  375. def switch_mode(mode='OBJECT', objects = []):
  376. active = None
  377. if objects:
  378. from bpy import context, ops
  379. active = objects[-1]
  380. context.view_layer.objects.active = active
  381. if (active):
  382. with context.temp_override(**{'active_object':active, 'selected_objects':objects}):
  383. ops.object.mode_set(mode=mode)
  384. return active
  385. # run this in Object mode, during bFinalize
  386. @preserve_active_object
  387. def bind_modifier_operator(modifier, operator):
  388. # now we have to bind it
  389. ob = modifier.id_data
  390. ob.modifiers.active = modifier
  391. import bpy
  392. bpy.context.view_layer.objects.active = ob
  393. # Context override does not do anything here... it isn't handled in the C code
  394. # I have verified this by building Blender with print statements to debug.
  395. # let's just make sure the target object has its modifiers disabled and update the dg
  396. targ_attr = "target"
  397. if hasattr(modifier, "object"): targ_attr = "object"
  398. target = getattr(modifier, targ_attr)
  399. if target:
  400. prWhite(f"Binding Deformer {modifier.name} to target {target.name}")
  401. operator(modifier=modifier.name)
  402. def get_default_collection(collection_type='WIDGET'):
  403. from ..preferences import get_bl_addon_object
  404. from bpy import data, context
  405. mantis_addon = get_bl_addon_object(raise_error=True)
  406. match collection_type:
  407. case "WIDGET":
  408. default_collection_name=mantis_addon.preferences.WidgetDefaultCollection
  409. case "CURVE":
  410. default_collection_name=mantis_addon.preferences.CurveDefaultCollection
  411. case "ARMATURE":
  412. default_collection_name=mantis_addon.preferences.MetaArmatureDefaultCollection
  413. if default_collection_name:
  414. if not (default_collection := data.collections.get(default_collection_name)):
  415. default_collection = data.collections.new(default_collection_name)
  416. context.scene.collection.children.link(default_collection)
  417. collection = default_collection
  418. else: collection = context.collection
  419. return collection
  420. def import_widget_obj(path,):
  421. from bpy.app import version as bpy_version
  422. from bpy import context, data
  423. from os import path as os_path
  424. file_name = os_path.split(path)[-1]
  425. obj_name = os_path.splitext(file_name)[0]
  426. collection = get_default_collection(collection_type='WIDGET')
  427. if bpy_version < (4,5,0):
  428. original_active = context.active_object
  429. # for blender versions prior to 4.5.0, we have to import with an operator
  430. from bpy.ops import wm as wm_ops
  431. ob_names_before = data.objects.keys()
  432. wm_ops.obj_import(
  433. filepath=path,
  434. check_existing=False,
  435. forward_axis='NEGATIVE_Z',
  436. up_axis='Y',
  437. validate_meshes=True,)
  438. # just make sure the active object doesn't change
  439. context.view_layer.objects.active = original_active
  440. # the below is a HACK... I can find the objects in the .obj file
  441. # by scanning the file for the "o" prefix and checking the name.
  442. # but that may be slow if the obj is big. which would make a bad widget!
  443. ob = None
  444. for ob in data.objects:
  445. if ob.name in ob_names_before: continue
  446. # this is easier than setting the active collection before import.
  447. for other_collection in ob.users_collection:
  448. other_collection.objects.unlink(ob)
  449. from math import pi as PI
  450. from mathutils import Matrix
  451. m = ob.data
  452. for v in m.vertices:
  453. v.co = Matrix.Rotation(PI/2, 4, 'X') @ v.co
  454. collection.objects.link(ob)
  455. return ob # return the first one, that should be the one
  456. else: # no new object was found - fail.
  457. # I don't expect this to happen unless there is an error in the operator.
  458. raise RuntimeError(f"Failed to import {file_name}. This is probably"
  459. "a bug or a corrupted file.")
  460. else:
  461. prWhite(f"INFO: using Geometry Nodes to import {file_name}")
  462. mesh = data.meshes.new(obj_name)
  463. ob = data.objects.new(name=obj_name, object_data=mesh)
  464. # we'll do a geometry nodes import
  465. collection.objects.link(ob)
  466. if (import_modifier := ob.modifiers.get("Import OBJ")) is None:
  467. import_modifier = ob.modifiers.new("Import OBJ", type='NODES')
  468. ng = data.node_groups.get("Import OBJ")
  469. if ng is None:
  470. from .geometry_node_graphgen import gen_import_obj_node_group
  471. ng = gen_import_obj_node_group()
  472. import_modifier.node_group = ng
  473. import_modifier["Socket_0"]=path
  474. return ob
  475. def import_object_from_file(path):
  476. # first let's check to see if we need it.
  477. from os import path as os_path
  478. file_name = os_path.split(path)[-1]
  479. obj_name = os_path.splitext(file_name)[0]
  480. extension = os_path.splitext(file_name)[1]
  481. if extension == '.obj':
  482. return import_widget_obj(path,)
  483. else:
  484. raise RuntimeError(f"Failed to parse filename {path}")
  485. def import_metarig_data(metarig_data : dict, ):
  486. from bpy import data, context
  487. from mathutils import Matrix
  488. from collections import deque
  489. # the metarig data is a dict with a bunch of nodes in it
  490. # start at node 'MANTIS_RESERVED'
  491. armature_data = metarig_data['MANTIS_RESERVED']
  492. children = deque(armature_data["children"].copy())
  493. if (armature := data.armatures.get(armature_data['name'])) is None:
  494. armature = data.armatures.new(armature_data['name'])
  495. # if we need to do anything here...
  496. if (armature_object := data.objects.get(armature_data['name'])) is None:
  497. armature_object = data.objects.new(armature_data['name'], object_data=armature)
  498. armature_object.matrix_world = Matrix(
  499. ( armature_data['matrix'][:4],
  500. armature_data['matrix'][4:8],
  501. armature_data['matrix'][8:12],
  502. armature_data['matrix'][12:16], )
  503. )
  504. prGreen (armature_data['name'])
  505. # have to add it to the view layer to switch modes.
  506. collection = get_default_collection(collection_type="ARMATURE")
  507. collection.objects.link(armature_object)
  508. # we'll do this to ensure it is actually in the scene for the mode switch
  509. context.scene.collection.objects.link(armature_object)
  510. switch_mode('EDIT', objects = [armature_object])
  511. while (children):
  512. child_name = children.pop()
  513. child_data = metarig_data[child_name]
  514. eb = armature.edit_bones.new(name=child_data['name'])
  515. if parent_name := child_data['parent']:
  516. eb.parent = armature.edit_bones[parent_name]
  517. eb.length = child_data['length']
  518. eb.matrix = Matrix(
  519. ( child_data['matrix'][:4],
  520. child_data['matrix'][4:8],
  521. child_data['matrix'][8:12],
  522. child_data['matrix'][12:16], )
  523. )
  524. displacement = eb.matrix.to_3x3().transposed().row[1] * child_data['length']
  525. eb.tail = eb.matrix.decompose()[0] + displacement
  526. children.extendleft (child_data['children'].copy())
  527. switch_mode('OBJECT', objects = [armature_object])
  528. # and now we can remove it from the scene collection, since it is in the Armature collection
  529. context.scene.collection.objects.unlink(armature_object)
  530. # note that this will not correct if the object exists and is wrong.
  531. return armature_object
  532. def import_curve_data_to_object(curve_name, curve_data):
  533. # the curve data will come as a single curve's data
  534. from bpy import data
  535. curve_object = data.objects.new(curve_name, data.curves.new(name=curve_name, type='CURVE'))
  536. curve_object.data.dimensions = '3D'
  537. prGreen (curve_name)
  538. for spline_data in curve_data:
  539. spline = curve_object.data.splines.new(type=spline_data['type'])
  540. points_data = spline_data['points']
  541. points_collection = spline.points
  542. if spline.type == 'BEZIER':
  543. # the points are bez_pts
  544. spline.bezier_points.add(len(points_data)-1)
  545. points_collection = spline.bezier_points
  546. else:
  547. spline.points.add(len(points_data)-1) # it starts with 1 already
  548. for i, point_data in enumerate(points_data):
  549. if spline.type == 'BEZIER':
  550. pt = spline.bezier_points[i]
  551. else:
  552. pt = spline.points[i]
  553. for prop in dir(pt):
  554. if prop == 'w':
  555. continue
  556. if prop == 'co' and spline.type != 'BEZIER':
  557. value = point_data[prop]
  558. pt.co = (value[0], value[1], value[2], point_data['w'])
  559. continue
  560. if prop in point_data.keys():
  561. setattr(pt, prop, point_data[prop])
  562. for prop in dir(spline):
  563. if prop in spline_data.keys():
  564. if prop in ['points', 'type', 'index']: continue
  565. setattr(spline, prop, spline_data[prop])
  566. collection = get_default_collection(collection_type="CURVE")
  567. collection.objects.link(curve_object)
  568. return curve_object
  569. def get_component_library_items(path='ADD_ARMATURE'):
  570. from os import path as os_path
  571. from ..preferences import get_bl_addon_object
  572. bl_mantis_addon = get_bl_addon_object()
  573. return_value=[]
  574. if bl_mantis_addon:
  575. match path:
  576. case 'ADD_ARMATURE':
  577. components_path = bl_mantis_addon.preferences.ComponentsLibraryFolder
  578. case 'AUTOLOAD':
  579. components_path = bl_mantis_addon.preferences.ComponentsAutoLoadFolder
  580. component_names = {}
  581. from os import walk as os_walk
  582. for path_root, dirs, files, in os_walk(components_path):
  583. for file in files:
  584. relative_file_name = os_path.join(os_path.sep.join(dirs), file)
  585. if file.endswith('.rig'):
  586. component_names[relative_file_name[:-4]] = relative_file_name
  587. if component_names.keys():
  588. for i, (name, path) in enumerate(component_names.items()):
  589. return_value.append( (path, name, path, 'NODE_TREE', i) )
  590. return return_value
  591. ##############################
  592. # READ TREE and also Schema Solve!
  593. ##############################
  594. # TODO: refactor the following two functions, they should be one function with arguments.
  595. def init_connections(mantis_node):
  596. c, hc = [], []
  597. for i in mantis_node.outputs.values():
  598. for l in i.links:
  599. # if l.from_node != mantis_node:
  600. # continue
  601. if l.is_hierarchy:
  602. hc.append(l.to_node)
  603. c.append(l.to_node)
  604. mantis_node.hierarchy_connections = hc
  605. mantis_node.connections = c
  606. def init_dependencies(mantis_node):
  607. c, hc = [], []
  608. for i in mantis_node.inputs.values():
  609. for l in i.links:
  610. # if l.to_node != mantis_node:
  611. # continue
  612. if l.is_hierarchy:
  613. hc.append(l.from_node)
  614. c.append(l.from_node)
  615. mantis_node.hierarchy_dependencies = hc
  616. mantis_node.dependencies = c
  617. def schema_dependency_handle_item(schema, all_mantis_nodes, item,):
  618. hierarchy = True
  619. from ..base_definitions import from_name_filter, to_name_filter
  620. if item.in_out == 'INPUT':
  621. dependencies = schema.dependencies
  622. hierarchy_dependencies = schema.hierarchy_dependencies
  623. parent_name = read_schema_type(item)
  624. if parent_name == 'Array':
  625. for schema_idname in ['SchemaArrayInput', 'SchemaArrayInputGet', 'SchemaArrayInputAll']:
  626. if (mantis_node := all_mantis_nodes.get( (*schema.signature, schema_idname) )):
  627. for to_link in mantis_node.outputs[item.name].links:
  628. if to_link.to_socket in to_name_filter:
  629. # hierarchy_reason='a'
  630. hierarchy = False
  631. for from_link in schema.inputs[item.identifier].links:
  632. if from_link.from_socket in from_name_filter:
  633. hierarchy = False
  634. # hierarchy_reason='b'
  635. if from_link.from_node not in dependencies:
  636. if hierarchy:
  637. hierarchy_dependencies.append(from_link.from_node)
  638. dependencies.append(from_link.from_node)
  639. if parent_name == 'Constant':
  640. if mantis_node := all_mantis_nodes.get((*schema.signature, 'SchemaConstInput')):
  641. for to_link in mantis_node.outputs[item.name].links:
  642. if to_link.to_socket in to_name_filter:
  643. # hierarchy_reason='dependencies'
  644. hierarchy = False
  645. for from_link in schema.inputs[item.identifier].links:
  646. if from_link.from_socket in from_name_filter:
  647. # hierarchy_reason='d'
  648. hierarchy = False
  649. if from_link.from_node not in dependencies:
  650. if hierarchy:
  651. hierarchy_dependencies.append(from_link.from_node)
  652. dependencies.append(from_link.from_node)
  653. if parent_name == 'Connection':
  654. if mantis_node := all_mantis_nodes.get((*schema.signature, 'SchemaIncomingConnection')):
  655. for to_link in mantis_node.outputs[item.name].links:
  656. if to_link.to_socket in to_name_filter:
  657. # hierarchy_reason='e'
  658. hierarchy = False
  659. for from_link in schema.inputs[item.identifier].links:
  660. if from_link.from_socket in from_name_filter:
  661. # hierarchy_reason='f'
  662. hierarchy = False
  663. if from_link.from_node not in dependencies:
  664. if hierarchy:
  665. hierarchy_dependencies.append(from_link.from_node)
  666. dependencies.append(from_link.from_node)
  667. def init_schema_dependencies(schema, all_mantis_nodes):
  668. """ Initialize the dependencies for Schema, and mark them as hierarchy or non-hierarchy dependencies
  669. Non-hierarchy dependencies are e.g. drivers and custom transforms.
  670. """
  671. tree = schema.ui_node.node_tree
  672. if tree is None:
  673. raise RuntimeError(f"Cannot get dependencies for schema {schema}")
  674. schema.dependencies = []
  675. schema.hierarchy_dependencies = []
  676. for l in schema.inputs["Schema Length"].links:
  677. schema.hierarchy_dependencies.append(l.from_node)
  678. if tree.interface:
  679. for item in tree.interface.items_tree:
  680. if item.item_type == 'PANEL':
  681. continue
  682. schema_dependency_handle_item(schema, all_mantis_nodes, item,)
  683. def check_and_add_root(n, roots, include_non_hierarchy=False):
  684. if (include_non_hierarchy * len(n.dependencies)) > 0:
  685. return
  686. elif len(n.hierarchy_dependencies) > 0:
  687. return
  688. roots.append(n)
  689. def get_link_in_out(link):
  690. from ..base_definitions import replace_types
  691. from_name, to_name = link.from_socket.node.name, link.to_socket.node.name
  692. # catch special bl_idnames and bunch the connections up
  693. if link.from_socket.node.bl_idname in replace_types:
  694. from_name = link.from_socket.node.bl_idname
  695. if link.to_socket.node.bl_idname in replace_types:
  696. to_name = link.to_socket.node.bl_idname
  697. return from_name, to_name
  698. def link_mantis_nodes(tree_path_names, link, local_mantis_nodes, from_suffix='', to_suffix=''):
  699. links = clear_reroutes([link])
  700. if links:
  701. link = links[0]
  702. else:
  703. return
  704. dummy_types = ["DUMMY", "DUMMY_SCHEMA"]
  705. from_name, to_name = get_link_in_out(link)
  706. from_mantis_node = local_mantis_nodes.get( (*tree_path_names, from_name+from_suffix) )
  707. to_mantis_node = local_mantis_nodes.get( (*tree_path_names, to_name+to_suffix))
  708. if (from_mantis_node and to_mantis_node):
  709. from_s, to_s = link.from_socket.name, link.to_socket.name
  710. if to_mantis_node.node_type in dummy_types: to_s = link.to_socket.identifier
  711. if from_mantis_node.node_type in dummy_types: from_s = link.from_socket.identifier
  712. try:
  713. connection = from_mantis_node.outputs[from_s].connect(node=to_mantis_node, socket=to_s, sort_id=link.multi_input_sort_id)
  714. if connection is None:
  715. prWhite(f"Already connected: {from_name}:{from_s}->{to_name}:{to_s}")
  716. return connection
  717. except KeyError as e:
  718. prRed(f"{from_mantis_node}:{from_s} or {to_mantis_node}:{to_s} missing; review the connections printed below:")
  719. print (from_mantis_node.outputs.keys())
  720. print (to_mantis_node.inputs.keys())
  721. raise e
  722. else:
  723. prRed(from_mantis_node, to_mantis_node, (*tree_path_names, from_name+from_suffix), (*tree_path_names, to_name+to_suffix))
  724. if to_mantis_node is None:
  725. prRed( "Error: mantis node with signature "
  726. f"{(*tree_path_names, to_name+to_suffix)} "
  727. "does not exist.")
  728. if from_mantis_node is None:
  729. prRed("Error: mantis node with signature "
  730. f"{(*tree_path_names, from_name+from_suffix)} "
  731. "does not exist.")
  732. raise RuntimeError(wrapRed(f"Link not connected: {from_mantis_node} -> {to_mantis_node} in tree" ))
  733. def get_all_dependencies(mantis_node):
  734. from ..base_definitions import GraphError
  735. """ find all dependencies for a mantis node"""
  736. nodes = []
  737. check_nodes = [mantis_node]
  738. nodes_checked = set()
  739. while (len(check_nodes) > 0):
  740. node = check_nodes.pop()
  741. nodes_checked.add (node)
  742. connected_nodes = node.hierarchy_dependencies
  743. for new_node in connected_nodes:
  744. if new_node in nodes:
  745. continue
  746. nodes.append(new_node)
  747. if new_node not in nodes_checked:
  748. check_nodes.append(new_node)
  749. return nodes
  750. def get_all_nodes_of_type(base_tree, bl_idname):
  751. nodes = []
  752. check_nodes = list(base_tree.nodes)
  753. while (len(check_nodes) > 0):
  754. node = check_nodes.pop()
  755. if node.bl_idname in bl_idname:
  756. nodes.append(node)
  757. if hasattr(node, "node_tree"):
  758. check_nodes.extend(list(node.node_tree.nodes))
  759. return nodes
  760. def trace_all_nodes_from_root(root, nodes):
  761. from ..base_definitions import GraphError
  762. """ find all dependencies for a mantis node"""
  763. nodes.add(root); check_nodes = [root]
  764. nodes_checked = set()
  765. while (len(check_nodes) > 0):
  766. node = check_nodes.pop(); nodes_checked.add (node)
  767. connected_nodes = []
  768. for output in node.outputs:
  769. for l in output.links:
  770. if l.to_node not in nodes:
  771. connected_nodes.append(l.to_node)
  772. for new_node in connected_nodes:
  773. nodes.add(new_node)
  774. if new_node not in nodes_checked:
  775. check_nodes.append(new_node)
  776. return nodes
  777. ##################################################################################################
  778. # misc
  779. ##################################################################################################
  780. # TODO: get the matrix to return a mathutils.Matrix so I don't need a function call here
  781. def to_mathutils_value(socket):
  782. if hasattr(socket, "default_value"):
  783. val = socket.default_value
  784. if socket.bl_idname in ['MatrixSocket']:
  785. return socket.TellValue()
  786. else:
  787. return val
  788. else:
  789. return None
  790. def all_trees_in_tree(base_tree, selected=False):
  791. """ Recursively finds all trees referenced in a given base-tree."""
  792. # note that this is recursive but not by tail-end recursion
  793. # a while-loop is a better way to do recursion in Python.
  794. trees = [base_tree]
  795. can_descend = True
  796. check_trees = [base_tree]
  797. while (len(check_trees) > 0): # this seems innefficient, why 2 loops?
  798. new_trees = []
  799. while (len(check_trees) > 0):
  800. tree = check_trees.pop()
  801. for node in tree.nodes:
  802. if selected == True and node.select == False:
  803. continue
  804. if new_tree := getattr(node, "node_tree", None):
  805. if new_tree in trees: continue
  806. new_trees.append(new_tree)
  807. trees.append(new_tree)
  808. check_trees = new_trees
  809. return trees
  810. # this is a destructive operation, not a pure function or whatever. That isn't good but I don't care.
  811. def SugiyamaGraph(tree, iterations):
  812. from grandalf.graphs import Vertex, Edge, Graph, graph_core
  813. class defaultview(object):
  814. w,h = 1,1
  815. xz = (0,0)
  816. graph = Graph()
  817. no_links = set()
  818. verts = {}
  819. for n in tree.nodes:
  820. if n.select == True:
  821. v = Vertex(n.name)
  822. v.view = defaultview()
  823. v.view.xy = n.location
  824. v.view.h = n.height*2.5
  825. v.view.w = n.width*2.2
  826. verts[n.name] = v
  827. no_links.add(n.name)
  828. graph.add_vertex(v)
  829. n.select=False
  830. edges = []
  831. inverted_edges=[]
  832. not_a_root = set()
  833. for link in tree.links:
  834. if (link.from_node.name not in verts.keys()) or (link.to_node.name not in verts.keys()):
  835. continue # problem??
  836. weight = 1 # maybe this is useful
  837. not_a_root.add(link.to_node.name) # if it has a edge-input it is not a root.
  838. e = Edge(verts[link.from_node.name], verts[link.to_node.name], weight)
  839. graph.add_edge(e)
  840. edges.append(e )
  841. if link.is_valid == False:
  842. inverted_edges.append(e)
  843. if link.from_node.name in no_links:
  844. no_links.remove(link.from_node.name)
  845. if link.to_node.name in no_links:
  846. no_links.remove(link.to_node.name)
  847. try:
  848. from grandalf.layouts import SugiyamaLayout
  849. # .C[0] is the first "graph core" that contains a connected graph.
  850. sug = SugiyamaLayout(graph.C[0])
  851. sug.init_all()
  852. sug.draw(iterations)
  853. # Digco is good for small graphs.
  854. # from grandalf.layouts import DigcoLayout
  855. # dco = DigcoLayout(graph.C[0])
  856. # dco.init_all()
  857. # dco.draw(iterations)
  858. except KeyboardInterrupt:
  859. pass # just use what it has calculated so far, I guess
  860. for v in graph.C[0].sV:
  861. for n in tree.nodes:
  862. if n.name == v.data:
  863. n.location.x = v.view.xy[1]
  864. n.location.y = v.view.xy[0]
  865. n.select = True
  866. # now we can take all the input nodes and try to put them in a sensible place
  867. # not sure why but this absolutely does not do anything
  868. for n_name in no_links:
  869. n = tree.nodes.get(n_name)
  870. next_node = None
  871. for output in n.outputs:
  872. if output.is_linked == True:
  873. next_node = output.links[0].to_node
  874. break
  875. # let's see if the next node
  876. if next_node:
  877. # need to find the other node in the same layer...
  878. other_node = None
  879. for s_input in next_node.inputs:
  880. if s_input.is_linked:
  881. other_node = s_input.links[0].from_node
  882. if other_node is n:
  883. continue
  884. else:
  885. break
  886. if other_node:
  887. n.location = other_node.location
  888. n.location.y -= other_node.height*2
  889. else: # we'll just position it next to the next node
  890. n.location = next_node.location
  891. n.location.x -= next_node.width*1.5
  892. def project_point_to_plane(point, origin, normal):
  893. return point - normal.dot(point- origin)*normal
  894. # If a $ character is in a value, it is an overlay request
  895. # in the root tree, supply it immediately, otherwise tag this node
  896. # once we finish pulling data from the tree we handle the overlay.
  897. def handle_overlay_detect(node, socket_id, value) -> str:
  898. mContext = node.mContext
  899. root_tree = mContext.base_tree
  900. root_items_tree = root_tree.interface.items_tree
  901. # detect if it is in the root tree, if so, handle immediately
  902. out_value = value
  903. if node.signature[0] is None and len(node.signature) ==2:
  904. for key, interface_item in root_items_tree.items():
  905. if interface_item.item_type != "SOCKET": continue
  906. if interface_item.in_out != "INPUT": continue
  907. # NOTE this will need to change if overlays extend to other data types
  908. if not hasattr(interface_item, "default_string"): continue
  909. out_value = out_value.replace("$"+key, interface_item.default_string)
  910. else: # since it is not root, we have to tag it.
  911. mContext.tagged_for_overlay.append((node.signature, socket_id))
  912. return out_value
  913. ####################################
  914. # CURVE STUFF
  915. ####################################
  916. def make_perpendicular(v1, v2):
  917. from ..base_definitions import FLOAT_EPSILON
  918. if (v1.length_squared < FLOAT_EPSILON) or (v2.length_squared < FLOAT_EPSILON):
  919. raise RuntimeError("Cannot generate perpendicular vetor for zero-length vector")
  920. projected = (v2.dot(v1) / v1.dot(v1)) * v1
  921. perpendicular = v2 - projected
  922. return perpendicular
  923. # this stuff could be branchless but I don't use it much TODO
  924. def cap(val, maxValue):
  925. if (val > maxValue):
  926. return maxValue
  927. return val
  928. def capMin(val, minValue):
  929. if (val < minValue):
  930. return minValue
  931. return val
  932. def wrap(min : float, max : float, value: float) -> float:
  933. range = max-min; remainder = value % range
  934. if remainder > max: return min + remainder-max
  935. else: return remainder
  936. def lerpVal(a, b, fac = 0.5):
  937. return a + ( (b-a) * fac)
  938. #wtf this doesn't do anything even remotely similar to wrap
  939. # HACK BAD FIXME UNBREAK ME BAD
  940. # I don't understand what this function does but I am using it in multiple places?
  941. def old_bad_wrap_that_should_be_refactored(val, maxValue, minValue = None):
  942. if (val > maxValue):
  943. return (-1 * ((maxValue - val) + 1))
  944. if ((minValue) and (val < minValue)):
  945. return (val + maxValue)
  946. return val
  947. #TODO clean this up
  948. def extract_spline_suffix(spline_index):
  949. return ".spline."+str(spline_index).zfill(3)+".extracted"
  950. def do_extract_spline(data, spline):
  951. remove_me = []
  952. for other_spline in data.splines:
  953. if other_spline != spline: remove_me.append(other_spline)
  954. while remove_me: data.splines.remove(remove_me.pop())
  955. def extract_spline(curve, spline_index):
  956. """ Given a curve object and spline index, returns a new object
  957. containing only the selcted spline. The new object is bound to
  958. the original curve.
  959. """
  960. if len(curve.data.splines) == 1:
  961. return curve # nothing to do here.
  962. spline_suffix = extract_spline_suffix(spline_index)
  963. from bpy import data
  964. if (new_ob := data.objects.get(curve.name+spline_suffix)) is None:
  965. new_ob=curve.copy(); new_ob.name=curve.name+spline_suffix
  966. # if the data exists, it is probably stale, so delete it and start over.
  967. if (old_data := data.objects.get(curve.data.name+spline_suffix)) is not None:
  968. data.curves.remove(old_data)
  969. new_data=curve.data.copy(); new_data.name=curve.data.name+spline_suffix
  970. new_ob.data = new_data
  971. # do not check for index error here, it is the calling function's responsibility
  972. do_extract_spline(new_data, new_data.splines[spline_index])
  973. return new_ob
  974. def bind_extracted_spline_to_curve(new_ob, curve):
  975. # Set up a relationship between the new object and the old object
  976. # now, weirdly enough - we can't use parenting very easily because Blender
  977. # defines the parent on a curve relative to the evaluated path animation
  978. # Setting the inverse matrix is too much work. Use Copy Transforms instead.
  979. from ..xForm_nodes import reset_object_data
  980. reset_object_data(new_ob)
  981. c = new_ob.constraints.new("COPY_TRANSFORMS"); c.target=curve
  982. new_ob.parent=curve
  983. return new_ob
  984. def get_extracted_spline_object(proto_curve, spline_index, mContext):
  985. # we're storing it separately like this to ensure all nodes use the same
  986. # object if they extract the same spline for use by Mantis.
  987. # this should be transparent to the user since it is working around a
  988. # a limitation in Blender.
  989. extracted_spline_name = proto_curve.name+extract_spline_suffix(spline_index)
  990. if curve := mContext.b_objects.get(extracted_spline_name):
  991. return curve
  992. else:
  993. curve = extract_spline(proto_curve, spline_index)
  994. if curve.name != proto_curve.name: # if there is only one spline, no
  995. bind_extracted_spline_to_curve(curve, proto_curve)# dupe is created.
  996. mContext.b_objects[extracted_spline_name] = curve
  997. return curve
  998. def nurbs_copy_bez_spline(curve, bez_spline, do_setup=True):
  999. other_spline= curve.data.splines.new('NURBS')
  1000. other_spline.use_endpoint_u=True
  1001. other_spline.use_bezier_u=True
  1002. bez_pts = bez_spline.bezier_points
  1003. bez_data=[]
  1004. for i, bez_pt in enumerate(bez_pts):
  1005. if i > 0:
  1006. bez_data.append(bez_pt.handle_left.copy())
  1007. bez_data.append(bez_pt.co.copy())
  1008. if i != len(bez_pts)-1:
  1009. bez_data.append(bez_pt.handle_right.copy())
  1010. other_spline.points.add(len(bez_data)-1)
  1011. for i, pt in enumerate(bez_data):
  1012. other_spline.points[i].co=(*pt,1.0) # add the W value here
  1013. if do_setup: # do the stuff that makes it behave the same as a bez spline
  1014. other_spline.use_endpoint_u = True; other_spline.use_bezier_u = True
  1015. other_spline.order_u=4 # set to 1 for poly
  1016. return other_spline
  1017. def RibbonMeshEdgeLengths(m, ribbon):
  1018. tE = ribbon[0]; bE = ribbon[1]; c = ribbon[2]
  1019. lengths = []
  1020. for i in range( len( tE ) ): #tE and bE are same length
  1021. if (c == True):
  1022. v1NextInd = tE[old_bad_wrap_that_should_be_refactored((i+1), len(tE) - 1)]
  1023. else:
  1024. v1NextInd = tE[cap((i+1) , len(tE) - 1 )]
  1025. v1 = m.vertices[tE[i]]; v1Next = m.vertices[v1NextInd]
  1026. if (c == True):
  1027. v2NextInd = bE[old_bad_wrap_that_should_be_refactored((i+1), len(bE) - 1)]
  1028. else:
  1029. v2NextInd = bE[cap((i+1) , len(bE) - 1 )]
  1030. v2 = m.vertices[bE[i]]; v2Next = m.vertices[v2NextInd]
  1031. v = v1.co.lerp(v2.co, 0.5); vNext = v1Next.co.lerp(v2Next.co, 0.5)
  1032. # get the center, edges may not be straight so total length
  1033. # of one edge may be more than the ribbon center's length
  1034. lengths.append(( v - vNext ).length)
  1035. return lengths
  1036. def EnsureCurveIsRibbon(crv, defaultRadius = 0.1):
  1037. from ..base_definitions import FLOAT_EPSILON
  1038. crvRadius = 0
  1039. crv.data.offset = 0
  1040. if (crv.data.bevel_depth < FLOAT_EPSILON):
  1041. crvRadius = crv.data.extrude
  1042. else: #Set ribbon from bevel depth
  1043. crvRadius = crv.data.bevel_depth
  1044. crv.data.bevel_depth = 0
  1045. crv.data.extrude = crvRadius
  1046. if (crvRadius < FLOAT_EPSILON):
  1047. crv.data.extrude = defaultRadius
  1048. def SetRibbonData(m, ribbon):
  1049. #maybe this could be incorporated into the DetectWireEdges function?
  1050. #maybe I can check for closed poly curves here? under what other circumstance
  1051. # will I find the ends of the wire have identical coordinates?
  1052. ribbonData = []
  1053. tE = ribbon[0].copy(); bE = ribbon[1].copy()# circle = ribbon[2]
  1054. #
  1055. lengths = RibbonMeshEdgeLengths(m, ribbon)
  1056. lengths.append(0)
  1057. totalLength = sum(lengths)
  1058. # m.calc_normals() #calculate normals
  1059. # it appears this has been removed.
  1060. for i, (t, b) in enumerate(zip(tE, bE)):
  1061. ind = old_bad_wrap_that_should_be_refactored( (i + 1), len(tE) - 1 )
  1062. tNext = tE[ind]; bNext = bE[ind]
  1063. ribbonData.append( ( (t,b), (tNext, bNext), lengths[i] ) )
  1064. #if this is a circle, the last v in vertData has a length, otherwise 0
  1065. return ribbonData, totalLength
  1066. def WireMeshEdgeLengths(m, wire):
  1067. circle = False
  1068. vIndex = wire.copy()
  1069. for e in m.edges:
  1070. if ((e.vertices[0] == vIndex[-1]) and (e.vertices[1] == vIndex[0])):
  1071. #this checks for an edge between the first and last vertex in the wire
  1072. circle = True
  1073. break
  1074. lengths = []
  1075. for i in range(len(vIndex)):
  1076. v = m.vertices[vIndex[i]]
  1077. if (circle == True):
  1078. vNextInd = vIndex[old_bad_wrap_that_should_be_refactored((i+1), len(vIndex) - 1)]
  1079. else:
  1080. vNextInd = vIndex[cap((i+1), len(vIndex) - 1 )]
  1081. vNext = m.vertices[vNextInd]
  1082. lengths.append(( v.co - vNext.co ).length)
  1083. #if this is a circular wire mesh, this should wrap instead of cap
  1084. return lengths
  1085. def GetDataFromWire(m, wire):
  1086. vertData = []
  1087. vIndex = wire.copy()
  1088. lengths = WireMeshEdgeLengths(m, wire)
  1089. lengths.append(0)
  1090. totalLength = sum(lengths)
  1091. for i, vInd in enumerate(vIndex):
  1092. #-1 to avoid IndexError
  1093. vNext = vIndex[ (old_bad_wrap_that_should_be_refactored(i+1, len(vIndex) - 1)) ]
  1094. vertData.append((vInd, vNext, lengths[i]))
  1095. #if this is a circle, the last v in vertData has a length, otherwise 0
  1096. return vertData, totalLength
  1097. def DetectWireEdges(mesh):
  1098. # Returns a list of vertex indices belonging to wire meshes
  1099. # NOTE: this assumes a mesh object with only wire meshes
  1100. ret = []
  1101. import bmesh
  1102. bm = bmesh.new()
  1103. try:
  1104. bm.from_mesh(mesh)
  1105. ends = []
  1106. for v in bm.verts:
  1107. if (len(v.link_edges) == 1):
  1108. ends.append(v.index)
  1109. for e in bm.edges:
  1110. assert (e.is_wire == True),"This function can only run on wire meshes"
  1111. if (e.verts[1].index - e.verts[0].index != 1):
  1112. ends.append(e.verts[1].index)
  1113. ends.append(e.verts[0].index)
  1114. for i in range(len(ends)//2): # // is floor division
  1115. beg = ends[i*2]
  1116. end = ends[(i*2)+1]
  1117. indices = [(j + beg) for j in range ((end - beg) + 1)]
  1118. ret.append(indices)
  1119. finally:
  1120. bm.free()
  1121. return ret
  1122. def FindNearestPointOnWireMesh(m, pointsList):
  1123. from mathutils import Vector
  1124. from mathutils.geometry import intersect_point_line
  1125. from math import sqrt
  1126. wires = DetectWireEdges(m)
  1127. ret = []
  1128. # prevFactor = None
  1129. for wire, points in zip(wires, pointsList):
  1130. vertData, total_length = GetDataFromWire(m, wire)
  1131. factorsOut = []
  1132. for p in points:
  1133. prevDist = float('inf')
  1134. curDist = float('inf')
  1135. v1 = None
  1136. v2 = None
  1137. for i in range(len(vertData) - 1):
  1138. #but it shouldn't check the last one
  1139. if (p == m.vertices[i].co):
  1140. v1 = vertData[i]
  1141. v2 = vertData[i+1]
  1142. offset = 0
  1143. break
  1144. else:
  1145. curDist = ( ((m.vertices[vertData[i][0]].co - p).length) +
  1146. ((m.vertices[vertData[i][1]].co - p).length) )/2
  1147. if (curDist < prevDist):
  1148. v1 = vertData[i]
  1149. v2 = vertData[i+1]
  1150. prevDist = curDist
  1151. offset = intersect_point_line(p, m.vertices[v1[0]].co,
  1152. m.vertices[v2[0]].co)[1]
  1153. if (offset < 0):
  1154. offset = 0
  1155. elif (offset > 1):
  1156. offset = 1
  1157. # Assume the vertices are in order
  1158. v1Length = 0
  1159. v2Length = v2[2]
  1160. for i in range(v1[0]):
  1161. v1Length += vertData[i][2]
  1162. factor = ((offset * (v2Length)) + v1Length )/total_length
  1163. factor = wrap(0, 1, factor) # doesn't hurt to wrap it if it's over 1 or less than 0
  1164. factorsOut.append(factor)
  1165. ret.append( factorsOut )
  1166. return ret
  1167. def mesh_from_curve(crv, context, ribbon=True):
  1168. """Utility function for converting a mesh to a curve
  1169. which will return the correct mesh even with modifiers"""
  1170. import bpy
  1171. m = None
  1172. bevel = crv.data.bevel_depth
  1173. extrude = crv.data.extrude
  1174. offset = crv.data.offset
  1175. try:
  1176. if (len(crv.modifiers) > 0):
  1177. do_unlink = False
  1178. if (not context.scene.collection.all_objects.get(crv.name)):
  1179. context.collection.objects.link(crv) # i guess this forces the dg to update it?
  1180. do_unlink = True
  1181. dg = context.view_layer.depsgraph
  1182. # just gonna modify it for now lol
  1183. if ribbon:
  1184. EnsureCurveIsRibbon(crv)
  1185. else:
  1186. crv.data.bevel_depth=0
  1187. crv.data.extrude=0
  1188. crv.data.offset=0
  1189. # try:
  1190. dg.update()
  1191. mOb = crv.evaluated_get(dg)
  1192. m = bpy.data.meshes.new_from_object(mOb)
  1193. m.name=crv.data.name+'_mesh'
  1194. if (do_unlink):
  1195. context.collection.objects.unlink(crv)
  1196. else: # (ಥ﹏ಥ) why can't I just use this !
  1197. # for now I will just do it like this
  1198. if ribbon:
  1199. EnsureCurveIsRibbon(crv)
  1200. else:
  1201. crv.data.bevel_depth=0
  1202. crv.data.extrude=0
  1203. crv.data.offset=0
  1204. m = bpy.data.meshes.new_from_object(crv)
  1205. finally:
  1206. crv.data.bevel_depth = bevel
  1207. crv.data.extrude = extrude
  1208. crv.data.offset = offset
  1209. return m
  1210. def DetectRibbon(f, bm, skipMe):
  1211. fFirst = f.index
  1212. cont = True
  1213. circle = False
  1214. tEdge, bEdge = [],[]
  1215. while (cont == True):
  1216. skipMe.add(f.index)
  1217. tEdge.append (f.loops[0].vert.index) # top-left
  1218. bEdge.append (f.loops[3].vert.index) # bottom-left
  1219. nEdge = bm.edges.get([f.loops[1].vert, f.loops[2].vert])
  1220. nFaces = nEdge.link_faces
  1221. if (len(nFaces) == 1):
  1222. cont = False
  1223. else:
  1224. for nFace in nFaces:
  1225. if (nFace != f):
  1226. f = nFace
  1227. break
  1228. if (f.index == fFirst):
  1229. cont = False
  1230. circle = True
  1231. if (cont == False): # we've reached the end, get the last two:
  1232. tEdge.append (f.loops[1].vert.index) # top-right
  1233. bEdge.append (f.loops[2].vert.index) # bottom-right
  1234. # this will create a loop for rings --
  1235. # "the first shall be the last and the last shall be first"
  1236. return (tEdge,bEdge,circle)
  1237. def DetectRibbons(m, fReport = None):
  1238. # Returns list of vertex indices belonging to ribbon mesh edges
  1239. # NOTE: this assumes a mesh object with only ribbon meshes
  1240. # ---DO NOT call this script with a mesh that isn't a ribbon!--- #
  1241. import bmesh
  1242. bm = bmesh.new()
  1243. bm.from_mesh(m)
  1244. mIslands, mIsland = [], []
  1245. skipMe = set()
  1246. bm.faces.ensure_lookup_table()
  1247. #first, get a list of mesh islands
  1248. for f in bm.faces:
  1249. if (f.index in skipMe):
  1250. continue #already done here
  1251. checkMe = [f]
  1252. while (len(checkMe) > 0):
  1253. facesFound = 0
  1254. for f in checkMe:
  1255. if (f.index in skipMe):
  1256. continue #already done here
  1257. mIsland.append(f)
  1258. skipMe.add(f.index)
  1259. for e in f.edges:
  1260. checkMe += e.link_faces
  1261. if (facesFound == 0):
  1262. #this is the last iteration
  1263. mIslands.append(mIsland)
  1264. checkMe, mIsland = [], []
  1265. ribbons = []
  1266. skipMe = set() # to store ends already checked
  1267. for mIsl in mIslands:
  1268. ribbon = None
  1269. first = float('inf')
  1270. for f in mIsl:
  1271. if (f.index in skipMe):
  1272. continue #already done here
  1273. if (f.index < first):
  1274. first = f.index
  1275. adjF = 0
  1276. for e in f.edges:
  1277. adjF+= (len(e.link_faces) - 1)
  1278. # every face other than this one is added to the list
  1279. if (adjF == 1):
  1280. ribbon = (DetectRibbon(f, bm, skipMe) )
  1281. break
  1282. if (ribbon == None):
  1283. ribbon = (DetectRibbon(bm.faces[first], bm, skipMe) )
  1284. ribbons.append(ribbon)
  1285. # print (ribbons)
  1286. return ribbons
  1287. def data_from_ribbon_mesh(m, factorsList, mat, ribbons = None, fReport = None):
  1288. #Note, factors list should be equal in length the the number of wires
  1289. #Now working for multiple wires, ugly tho
  1290. if (ribbons == None):
  1291. ribbons = DetectRibbons(m, fReport=fReport)
  1292. if (ribbons is None):
  1293. if (fReport):
  1294. fReport(type = {'ERROR'}, message="No ribbon to get data from.")
  1295. else:
  1296. print ("No ribbon to get data from.")
  1297. return None
  1298. ret = []
  1299. for factors, ribbon in zip(factorsList, ribbons):
  1300. points = []
  1301. widths = []
  1302. normals = []
  1303. ribbonData, totalLength = SetRibbonData(m, ribbon)
  1304. for fac in factors:
  1305. if (fac == 0):
  1306. data = ribbonData[0]
  1307. curFac = 0
  1308. elif (fac == 1):
  1309. data = ribbonData[-1]
  1310. curFac = 0
  1311. else:
  1312. targetLength = totalLength * fac
  1313. data = ribbonData[0]
  1314. curLength = 0
  1315. for ( (t, b), (tNext, bNext), length,) in ribbonData:
  1316. if (curLength >= targetLength):
  1317. break
  1318. curLength += length
  1319. data = ( (t, b), (tNext, bNext), length,)
  1320. targetLengthAtEdge = (curLength - targetLength)
  1321. if (targetLength == 0):
  1322. curFac = 0
  1323. elif (targetLength == totalLength):
  1324. curFac = 1
  1325. else:
  1326. # NOTE: This can be Zero. Find out why!
  1327. if data[2] == 0:
  1328. curFac=0
  1329. else:
  1330. curFac = 1 - (targetLengthAtEdge/ data[2]) #length
  1331. t1 = m.vertices[data[0][0]]; b1 = m.vertices[data[0][1]]
  1332. t2 = m.vertices[data[1][0]]; b2 = m.vertices[data[1][1]]
  1333. #location
  1334. loc1 = (t1.co).lerp(b1.co, 0.5)
  1335. loc2 = (t2.co).lerp(b2.co, 0.5)
  1336. #width
  1337. w1 = (t1.co - b1.co).length/2
  1338. w2 = (t2.co - b2.co).length/2 #radius, not diameter
  1339. #normal
  1340. n1 = (t1.normal).slerp(b1.normal, 0.5)
  1341. n2 = (t1.normal).slerp(b2.normal, 0.5)
  1342. if ((data[0][0] > data[1][0]) and (ribbon[2] == False)):
  1343. curFac = 0
  1344. #don't interpolate if at the end of a ribbon that isn't circular
  1345. if ( 0 < curFac < 1):
  1346. outPoint = loc1.lerp(loc2, curFac)
  1347. outNorm = n1.lerp(n2, curFac)
  1348. outWidth = w1 + ( (w2-w1) * curFac)
  1349. elif (curFac <= 0):
  1350. outPoint = loc1.copy()
  1351. outNorm = n1
  1352. outWidth = w1
  1353. elif (curFac >= 1):
  1354. outPoint = loc2.copy()
  1355. outNorm = n2
  1356. outWidth = w2
  1357. outPoint = mat @ outPoint
  1358. outNorm.normalize()
  1359. points.append ( outPoint.copy() ) #copy because this is an actual vertex location
  1360. widths.append ( outWidth )
  1361. normals.append( outNorm )
  1362. ret.append( (points, widths, normals) )
  1363. return ret # this is a list of tuples containing three lists
  1364. #This bisection search is generic, and it searches based on the
  1365. # magnitude of the error, rather than the sign.
  1366. # If the sign of the error is meaningful, a simpler function
  1367. # can be used.
  1368. def do_bisect_search_by_magnitude(
  1369. owner,
  1370. attribute,
  1371. index = None,
  1372. test_function = None,
  1373. modify = None,
  1374. max_iterations = 10000,
  1375. threshold = 0.0001,
  1376. thresh2 = 0.0005,
  1377. context = None,
  1378. update_dg = None,
  1379. ):
  1380. from math import floor
  1381. i = 0; best_so_far = 0; best = float('inf')
  1382. min = 0; center = max_iterations//2; max = max_iterations
  1383. # enforce getting the absolute value, in case the function has sign information
  1384. # The sign may be useful in a sign-aware bisect search, but this one is more robust!
  1385. test = lambda : abs(test_function(owner, attribute, index, context = context,))
  1386. while (i <= max_iterations):
  1387. upper = (max - ((max-center))//2)
  1388. modify(owner, attribute, index, upper, context = context); error1 = test()
  1389. lower = (center - ((center-min))//2)
  1390. modify(owner, attribute, index, lower, context = context); error2 = test()
  1391. if (error1 < error2):
  1392. min = center
  1393. center, check = upper, upper
  1394. error = error1
  1395. else:
  1396. max = center
  1397. center, check = lower, lower
  1398. error = error2
  1399. if (error <= threshold) or (min == max-1):
  1400. break
  1401. if (error < thresh2):
  1402. j = min
  1403. while (j < max):
  1404. modify(owner, attribute, index, j * 1/max_iterations, context = context)
  1405. error = test()
  1406. if (error < best):
  1407. best_so_far = j; best = error
  1408. if (error <= threshold):
  1409. break
  1410. j+=1
  1411. else: # loop has completed without finding a solution
  1412. i = best_so_far; error = test()
  1413. modify(owner, attribute, index, best_so_far, context = context)
  1414. break
  1415. if (error < best):
  1416. best_so_far = check; best = error
  1417. i+=1
  1418. if update_dg:
  1419. update_dg.update()
  1420. else: # Loop has completed without finding a solution
  1421. i = best_so_far
  1422. modify(owner, attribute, best_so_far, context = context); i+=1