utilities.py 63 KB

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