utilities.py 46 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093
  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. # A fuction for getting to the end of a Reroute.
  19. # TODO: this seems really inefficient!
  20. def socket_seek(start_link, links):
  21. link = start_link
  22. while(link.from_socket):
  23. for newlink in links:
  24. if link.from_socket.node.inputs:
  25. if newlink.to_socket == link.from_socket.node.inputs[0]:
  26. link=newlink; break
  27. else:
  28. break
  29. return link.from_socket
  30. # this creates fake links that have the same interface as Blender's
  31. # so that I can bypass Reroutes
  32. def clear_reroutes(links):
  33. from .base_definitions import DummyLink
  34. kept_links, rerouted_starts = [], []
  35. rerouted = []
  36. all_links = links.copy()
  37. while(all_links):
  38. link = all_links.pop()
  39. to_cls = link.to_socket.node.bl_idname
  40. from_cls = link.from_socket.node.bl_idname
  41. reroute_classes = ["NodeReroute"]
  42. if (to_cls in reroute_classes and
  43. from_cls in reroute_classes):
  44. rerouted.append(link)
  45. elif (to_cls in reroute_classes and not
  46. from_cls in reroute_classes):
  47. rerouted.append(link)
  48. elif (from_cls in reroute_classes and not
  49. to_cls in reroute_classes):
  50. rerouted_starts.append(link)
  51. else:
  52. kept_links.append(link)
  53. for start in rerouted_starts:
  54. from_socket = socket_seek(start, rerouted)
  55. 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 )
  56. kept_links.append(new_link)
  57. return kept_links
  58. def tree_from_nc(sig, base_tree):
  59. if (sig[0] == 'MANTIS_AUTOGENERATED'):
  60. sig = sig[:-2] # cut off the end part of the signature (because it uses socket.name and socket.identifier)
  61. # this will lead to totally untraceble bugs in the event of a change in how signatures are assigned
  62. tree = base_tree
  63. for i, path_item in enumerate(sig):
  64. if (i == 0) or (i == len(sig) - 1):
  65. continue
  66. tree = tree.nodes.get(path_item).node_tree
  67. return tree
  68. def get_node_prototype(sig, base_tree):
  69. return tree_from_nc(sig, base_tree).nodes.get( sig[-1] )
  70. ##################################################################################################
  71. # groups and changing sockets -- this is used extensively by Schema.
  72. ##################################################################################################
  73. # this one returns None if there is an error.
  74. def get_socket_maps(node, force=False):
  75. maps = [{}, {}]
  76. node_collection = ["inputs", "outputs"]
  77. links = ["from_socket", "to_socket"]
  78. for collection, map, linked_socket in zip(node_collection, maps, links):
  79. for sock in getattr(node, collection):
  80. if sock.is_linked:
  81. other_sockets = []
  82. # Sort the links first (in case they are mult-input), because Blender doesn't
  83. links = sorted(list(sock.links), key = lambda l : l.multi_input_sort_id)
  84. # HACK here because Blender will crash if the socket values in the NodeReroute
  85. # are mutated. Because this seems to happen in a deffered way, I can't account
  86. # for it except by checking the node later...
  87. # TODO: The fact that I need this hack means I can probably solve this problem
  88. # for all node types in a safer way, since they may also be dynamic somehow
  89. for l in links:
  90. if "from" in linked_socket and l.from_node.bl_idname == "NodeReroute":
  91. other_sockets.append(l.from_node)
  92. elif "to" in linked_socket and l.to_node.bl_idname == "NodeReroute":
  93. other_sockets.append(l.to_node)
  94. else:
  95. other_sockets.append(getattr(l, linked_socket))
  96. map[sock.identifier]= other_sockets
  97. elif hasattr(sock, "default_value"):
  98. if sock.get("default_value") is not None:
  99. val = sock['default_value']
  100. elif sock.bl_idname == "EnumCurveSocket" and sock.get("default_value") is None:
  101. # HACK I need to add this special case because during file-load,
  102. # this value is None and should not be altered until it is set once.
  103. continue
  104. elif (val := sock.default_value) is not None:
  105. pass
  106. elif not force:
  107. continue
  108. map[sock.identifier]=val
  109. else:
  110. from .socket_definitions import no_default_value
  111. if sock.bl_idname in no_default_value:
  112. map[sock.identifier]=None
  113. else:
  114. raise RuntimeError(f"ERROR: Could not get socket data for socket of type: {sock.bl_idname}")
  115. return maps
  116. # this function is completely overloaded with different purposes and code paths
  117. # TODO refactor everything that funnels into this function
  118. # make this stuff simpler.
  119. def do_relink(node, s, map, in_out='INPUT', parent_name = ''):
  120. if not node.__class__.is_registered_node_type(): return
  121. tree = node.id_data; interface_in_out = 'OUTPUT' if in_out == 'INPUT' else 'INPUT'
  122. if hasattr(node, "node_tree"):
  123. tree = node.node_tree
  124. interface_in_out=in_out
  125. from bpy.types import NodeSocket, Node
  126. get_string = '__extend__'
  127. if s: get_string = s.identifier
  128. from .base_definitions import SchemaUINode
  129. if (hasattr(node, "node_tree") or isinstance(node, SchemaUINode)) and get_string not in map.keys():
  130. # this happens when we are creating a new node group and need to update it from nothing.
  131. return
  132. val = map[get_string] # this will throw an error if the socket isn't there. Good!
  133. if isinstance(val, list):
  134. for sub_val in val:
  135. # this will only happen once because it assigns s, so it is safe to do in the for loop.
  136. if s is None:
  137. name = unique_socket_name(node, sub_val, tree)
  138. sock_type = sub_val.bl_idname
  139. if parent_name:
  140. interface_socket = update_interface(tree.interface, name, interface_in_out, sock_type, parent_name)
  141. if in_out =='INPUT':
  142. s = node.inputs.new(sock_type, name, identifier=interface_socket.identifier)
  143. else:
  144. s = node.outputs.new(sock_type, name, identifier=interface_socket.identifier)
  145. if parent_name == 'Array': s.display_shape='SQUARE_DOT'
  146. # then move it up and delete the other link.
  147. # this also needs to modify the interface of the node tree.
  148. if isinstance(sub_val, NodeSocket):
  149. l = None
  150. if in_out =='INPUT':
  151. l = node.id_data.links.new(input=sub_val, output=s)
  152. else:
  153. l = node.id_data.links.new(input=s, output=sub_val)
  154. if l is None:
  155. raise RuntimeError("Could not create link")
  156. elif isinstance(sub_val, Node):
  157. l = None
  158. # this happens when it is a NodeReroute
  159. if not s.is_output:
  160. l = node.id_data.links.new(input=sub_val.outputs[0], output=s)
  161. else:
  162. l = node.id_data.links.new(input=s, output=sub_val.inputs[0])
  163. if l is None:
  164. raise RuntimeError("Could not create link")
  165. else:
  166. raise RuntimeError("Unhandled case in do_relink()")
  167. elif get_string != "__extend__":
  168. if not s.is_output:
  169. try:
  170. s.default_value = val
  171. except (AttributeError, ValueError): # must be readonly or maybe it doesn't have a d.v.
  172. pass
  173. def update_interface(interface, name, in_out, sock_type, parent_name):
  174. if parent_name:
  175. if not (interface_parent := interface.items_tree.get(parent_name)):
  176. interface_parent = interface.new_panel(name=parent_name)
  177. socket = interface.new_socket(name=name,in_out=in_out, socket_type=sock_type, parent=interface_parent)
  178. if parent_name == 'Connection':
  179. in_out = 'OUTPUT' if in_out == 'INPUT' else 'INPUT' # flip this make sure connections always do both
  180. interface.new_socket(name=name,in_out=in_out, socket_type=sock_type, parent=interface_parent)
  181. return socket
  182. else:
  183. raise RuntimeError(wrapRed("Cannot add interface item to tree without specifying type."))
  184. #UGLY BAD REFACTOR
  185. def relink_socket_map_add_socket(node, socket_collection, item, in_out=None,):
  186. if not in_out: in_out=item.in_out
  187. if node.bl_idname in ['MantisSchemaGroup'] and item.parent and item.parent.name == 'Array':
  188. multi = True if in_out == 'INPUT' else False
  189. s = socket_collection.new(type=item.socket_type, name=item.name, identifier=item.identifier, use_multi_input=multi)
  190. else:
  191. s = socket_collection.new(type=item.socket_type, name=item.name, identifier=item.identifier)
  192. if item.parent.name == 'Array': s.display_shape = 'SQUARE_DOT'
  193. return s
  194. # TODO REFACTOR THIS
  195. # I did this awful thing because I needed the above code
  196. # but I have provided this interface to Mantis
  197. # I did not follow the Single Responsibility Principle
  198. # I am now suffering for it, as I rightly deserve.
  199. def relink_socket_map(node, socket_collection, map, item, in_out=None,):
  200. s = relink_socket_map_add_socket(node, socket_collection, item, in_out=None,)
  201. do_relink(node, s, map)
  202. def unique_socket_name(node, other_socket, tree):
  203. name_stem = other_socket.bl_label; num=0
  204. # if hasattr(other_socket, "default_value"):
  205. # name_stem = type(other_socket.default_value).__name__
  206. for item in tree.interface.items_tree:
  207. if item.item_type == 'PANEL': continue
  208. if other_socket.is_output and item.in_out == 'INPUT': continue
  209. if not other_socket.is_output and item.in_out == 'OUTPUT': continue
  210. if name_stem in item.name: num+=1
  211. name = name_stem + '.' + str(num).zfill(3)
  212. return name
  213. ##############################
  214. # READ TREE and also Schema Solve!
  215. ##############################
  216. # TODO: refactor the following two functions, they should be one function with arguments.
  217. def init_connections(nc):
  218. c, hc = [], []
  219. for i in nc.outputs.values():
  220. for l in i.links:
  221. # if l.from_node != nc:
  222. # continue
  223. if l.is_hierarchy:
  224. hc.append(l.to_node)
  225. c.append(l.to_node)
  226. nc.hierarchy_connections = hc
  227. nc.connections = c
  228. def init_dependencies(nc):
  229. c, hc = [], []
  230. for i in nc.inputs.values():
  231. for l in i.links:
  232. # if l.to_node != nc:
  233. # continue
  234. if l.is_hierarchy:
  235. hc.append(l.from_node)
  236. c.append(l.from_node)
  237. nc.hierarchy_dependencies = hc
  238. nc.dependencies = c
  239. def schema_dependency_handle_item(schema, all_nc, item,):
  240. hierarchy = True
  241. from .base_definitions import from_name_filter, to_name_filter
  242. if item.in_out == 'INPUT':
  243. dependencies = schema.dependencies
  244. hierarchy_dependencies = schema.hierarchy_dependencies
  245. if item.parent and item.parent.name == 'Array':
  246. for schema_idname in ['SchemaArrayInput', 'SchemaArrayInputGet', 'SchemaArrayInputAll']:
  247. if (nc := all_nc.get( (*schema.signature, schema_idname) )):
  248. for to_link in nc.outputs[item.name].links:
  249. if to_link.to_socket in to_name_filter:
  250. # hierarchy_reason='a'
  251. hierarchy = False
  252. for from_link in schema.inputs[item.identifier].links:
  253. if from_link.from_socket in from_name_filter:
  254. hierarchy = False
  255. # hierarchy_reason='b'
  256. if from_link.from_node not in dependencies:
  257. if hierarchy:
  258. hierarchy_dependencies.append(from_link.from_node)
  259. dependencies.append(from_link.from_node)
  260. if item.parent and item.parent.name == 'Constant':
  261. if nc := all_nc.get((*schema.signature, 'SchemaConstInput')):
  262. for to_link in nc.outputs[item.name].links:
  263. if to_link.to_socket in to_name_filter:
  264. # hierarchy_reason='dependencies'
  265. hierarchy = False
  266. for from_link in schema.inputs[item.identifier].links:
  267. if from_link.from_socket in from_name_filter:
  268. # hierarchy_reason='d'
  269. hierarchy = False
  270. if from_link.from_node not in dependencies:
  271. if hierarchy:
  272. hierarchy_dependencies.append(from_link.from_node)
  273. dependencies.append(from_link.from_node)
  274. if item.parent and item.parent.name == 'Connection':
  275. if nc := all_nc.get((*schema.signature, 'SchemaIncomingConnection')):
  276. for to_link in nc.outputs[item.name].links:
  277. if to_link.to_socket in to_name_filter:
  278. # hierarchy_reason='e'
  279. hierarchy = False
  280. for from_link in schema.inputs[item.identifier].links:
  281. if from_link.from_socket in from_name_filter:
  282. # hierarchy_reason='f'
  283. hierarchy = False
  284. if from_link.from_node not in dependencies:
  285. if hierarchy:
  286. hierarchy_dependencies.append(from_link.from_node)
  287. dependencies.append(from_link.from_node)
  288. def init_schema_dependencies(schema, all_nc):
  289. """ Initialize the dependencies for Schema, and mark them as hierarchy or non-hierarchy dependencies
  290. Non-hierarchy dependencies are e.g. drivers and custom transforms.
  291. """
  292. tree = schema.prototype.node_tree
  293. if tree is None:
  294. raise RuntimeError(f"Cannot get dependencies for schema {schema}")
  295. schema.dependencies = []
  296. schema.hierarchy_dependencies = []
  297. for l in schema.inputs["Schema Length"].links:
  298. schema.hierarchy_dependencies.append(l.from_node)
  299. if tree.interface:
  300. for item in tree.interface.items_tree:
  301. if item.item_type == 'PANEL':
  302. continue
  303. schema_dependency_handle_item(schema, all_nc, item,)
  304. def check_and_add_root(n, roots, include_non_hierarchy=False):
  305. if (include_non_hierarchy * len(n.dependencies)) > 0:
  306. return
  307. elif len(n.hierarchy_dependencies) > 0:
  308. return
  309. roots.append(n)
  310. def get_link_in_out(link):
  311. from .base_definitions import replace_types
  312. from_name, to_name = link.from_socket.node.name, link.to_socket.node.name
  313. # catch special bl_idnames and bunch the connections up
  314. if link.from_socket.node.bl_idname in replace_types:
  315. from_name = link.from_socket.node.bl_idname
  316. if link.to_socket.node.bl_idname in replace_types:
  317. to_name = link.to_socket.node.bl_idname
  318. return from_name, to_name
  319. def link_node_containers(tree_path_names, link, local_nc, from_suffix='', to_suffix=''):
  320. dummy_types = ["DUMMY", "DUMMY_SCHEMA"]
  321. from_name, to_name = get_link_in_out(link)
  322. nc_from = local_nc.get( (*tree_path_names, from_name+from_suffix) )
  323. nc_to = local_nc.get( (*tree_path_names, to_name+to_suffix))
  324. if (nc_from and nc_to):
  325. from_s, to_s = link.from_socket.name, link.to_socket.name
  326. if nc_to.node_type in dummy_types: to_s = link.to_socket.identifier
  327. if nc_from.node_type in dummy_types: from_s = link.from_socket.identifier
  328. try:
  329. connection = nc_from.outputs[from_s].connect(node=nc_to, socket=to_s, sort_id=link.multi_input_sort_id)
  330. if connection is None:
  331. prWhite(f"Already connected: {from_name}:{from_s}->{to_name}:{to_s}")
  332. return connection
  333. except KeyError as e:
  334. prRed(f"{nc_from}:{from_s} or {nc_to}:{to_s} missing; review the connections printed below:")
  335. print (nc_from.outputs.keys())
  336. print (nc_to.inputs.keys())
  337. raise e
  338. else:
  339. prRed(nc_from, nc_to, (*tree_path_names, from_name+from_suffix), (*tree_path_names, to_name+to_suffix))
  340. raise RuntimeError(wrapRed("Link not connected: %s -> %s in tree %s" % (from_name, to_name, tree_path_names[-1])))
  341. def get_all_dependencies(nc):
  342. from .base_definitions import GraphError
  343. """ find all dependencies for a mantis node"""
  344. nodes = []
  345. check_nodes = [nc]
  346. nodes_checked = set()
  347. while (len(check_nodes) > 0):
  348. node = check_nodes.pop()
  349. nodes_checked.add (node)
  350. connected_nodes = node.hierarchy_dependencies
  351. for new_node in connected_nodes:
  352. if new_node in nodes:
  353. continue
  354. nodes.append(new_node)
  355. if new_node not in nodes_checked:
  356. check_nodes.append(new_node)
  357. return nodes
  358. def get_all_nodes_of_type(base_tree, bl_idname):
  359. nodes = []
  360. check_nodes = list(base_tree.nodes)
  361. while (len(check_nodes) > 0):
  362. node = check_nodes.pop()
  363. if node.bl_idname in bl_idname:
  364. nodes.append(node)
  365. if hasattr(node, "node_tree"):
  366. check_nodes.extend(list(node.node_tree.nodes))
  367. return nodes
  368. ##################################################################################################
  369. # misc
  370. ##################################################################################################
  371. # TODO: get the matrix to return a mathutils.Matrix so I don't need a function call here
  372. def to_mathutils_value(socket):
  373. if hasattr(socket, "default_value"):
  374. val = socket.default_value
  375. if socket.bl_idname in ['MatrixSocket']:
  376. return socket.TellValue()
  377. else:
  378. return val
  379. else:
  380. return None
  381. def all_trees_in_tree(base_tree, selected=False):
  382. """ Recursively finds all trees referenced in a given base-tree."""
  383. # note that this is recursive but not by tail-end recursion
  384. # a while-loop is a better way to do recursion in Python.
  385. trees = [base_tree]
  386. can_descend = True
  387. check_trees = [base_tree]
  388. while (len(check_trees) > 0): # this seems innefficient, why 2 loops?
  389. new_trees = []
  390. while (len(check_trees) > 0):
  391. tree = check_trees.pop()
  392. for node in tree.nodes:
  393. if selected == True and node.select == False:
  394. continue
  395. if new_tree := getattr(node, "node_tree", None):
  396. if new_tree in trees: continue
  397. new_trees.append(new_tree)
  398. trees.append(new_tree)
  399. check_trees = new_trees
  400. return trees
  401. # this is a destructive operation, not a pure function or whatever. That isn't good but I don't care.
  402. def SugiyamaGraph(tree, iterations):
  403. from grandalf.graphs import Vertex, Edge, Graph, graph_core
  404. class defaultview(object):
  405. w,h = 1,1
  406. xz = (0,0)
  407. no_links = set()
  408. verts = {}
  409. for n in tree.nodes:
  410. has_links=False
  411. for inp in n.inputs:
  412. if inp.is_linked:
  413. has_links=True
  414. break
  415. else:
  416. no_links.add(n.name)
  417. for out in n.outputs:
  418. if out.is_linked:
  419. has_links=True
  420. break
  421. else:
  422. try:
  423. no_links.remove(n.name)
  424. except KeyError:
  425. pass
  426. if not has_links:
  427. continue
  428. v = Vertex(n.name)
  429. v.view = defaultview()
  430. v.view.xy = n.location
  431. v.view.h = n.height*2.5
  432. v.view.w = n.width*2.2
  433. verts[n.name] = v
  434. edges = []
  435. for link in tree.links:
  436. weight = 1 # maybe this is useful
  437. edges.append(Edge(verts[link.from_node.name], verts[link.to_node.name], weight) )
  438. graph = Graph(verts.values(), edges)
  439. from grandalf.layouts import SugiyamaLayout
  440. sug = SugiyamaLayout(graph.C[0]) # no idea what .C[0] is
  441. roots=[]
  442. for node in tree.nodes:
  443. has_links=False
  444. for inp in node.inputs:
  445. if inp.is_linked:
  446. has_links=True
  447. break
  448. for out in node.outputs:
  449. if out.is_linked:
  450. has_links=True
  451. break
  452. if not has_links:
  453. continue
  454. if len(node.inputs)==0:
  455. roots.append(verts[node.name])
  456. else:
  457. for inp in node.inputs:
  458. if inp.is_linked==True:
  459. break
  460. else:
  461. roots.append(verts[node.name])
  462. sug.init_all(roots=roots,)
  463. sug.draw(iterations)
  464. for v in graph.C[0].sV:
  465. for n in tree.nodes:
  466. if n.name == v.data:
  467. n.location.x = v.view.xy[1]
  468. n.location.y = v.view.xy[0]
  469. # now we can take all the input nodes and try to put them in a sensible place
  470. for n_name in no_links:
  471. n = tree.nodes.get(n_name)
  472. next_n = None
  473. next_node = None
  474. for output in n.outputs:
  475. if output.is_linked == True:
  476. next_node = output.links[0].to_node
  477. break
  478. # let's see if the next node
  479. if next_node:
  480. # need to find the other node in the same layer...
  481. other_node = None
  482. for s_input in next_node.inputs:
  483. if s_input.is_linked:
  484. other_node = s_input.links[0].from_node
  485. if other_node is n:
  486. continue
  487. else:
  488. break
  489. if other_node:
  490. n.location = other_node.location
  491. n.location.y -= other_node.height*2
  492. else: # we'll just position it next to the next node
  493. n.location = next_node.location
  494. n.location.x -= next_node.width*1.5
  495. def project_point_to_plane(point, origin, normal):
  496. return point - normal.dot(point- origin)*normal
  497. ##################################################################################################
  498. # stuff I should probably refactor!!
  499. ##################################################################################################
  500. # This is really, really stupid way to do this
  501. def gen_nc_input_for_data(socket):
  502. # Class List #TODO deduplicate
  503. from . import xForm_containers, link_containers, misc_nodes, primitives_containers, deformer_containers, math_containers, schema_containers
  504. from .internal_containers import NoOpNode
  505. classes = {}
  506. for module in [xForm_containers, link_containers, misc_nodes, primitives_containers, deformer_containers, math_containers, schema_containers]:
  507. for cls in module.TellClasses():
  508. classes[cls.__name__] = cls
  509. #
  510. socket_class_map = {
  511. "MatrixSocket" : classes["InputMatrix"],
  512. "xFormSocket" : None,
  513. "RelationshipSocket" : NoOpNode,
  514. "DeformerSocket" : NoOpNode,
  515. "GeometrySocket" : classes["InputExistingGeometryData"],
  516. "EnableSocket" : classes["InputBoolean"],
  517. "HideSocket" : classes["InputBoolean"],
  518. #
  519. "DriverSocket" : None,
  520. "DriverVariableSocket" : None,
  521. "FCurveSocket" : None,
  522. "KeyframeSocket" : None,
  523. "BoneCollectionSocket" : classes["InputString"],
  524. #
  525. "xFormParameterSocket" : None,
  526. "ParameterBoolSocket" : classes["InputBoolean"],
  527. "ParameterIntSocket" : classes["InputFloat"], #TODO: make an Int node for this
  528. "ParameterFloatSocket" : classes["InputFloat"],
  529. "ParameterVectorSocket" : classes["InputVector"],
  530. "ParameterStringSocket" : classes["InputString"],
  531. #
  532. "TransformSpaceSocket" : classes["InputTransformSpace"],
  533. "BooleanSocket" : classes["InputBoolean"],
  534. "BooleanThreeTupleSocket" : classes["InputBooleanThreeTuple"],
  535. "RotationOrderSocket" : classes["InputRotationOrder"],
  536. "QuaternionSocket" : None,
  537. "QuaternionSocketAA" : None,
  538. "UnsignedIntSocket" : classes["InputFloat"],
  539. "IntSocket" : classes["InputFloat"],
  540. "StringSocket" : classes["InputString"],
  541. #
  542. "BoolUpdateParentNode" : classes["InputBoolean"],
  543. "IKChainLengthSocket" : classes["InputFloat"],
  544. "EnumInheritScale" : classes["InputString"],
  545. "EnumRotationMix" : classes["InputString"],
  546. "EnumRotationMixCopyTransforms" : classes["InputString"],
  547. "EnumMaintainVolumeStretchTo" : classes["InputString"],
  548. "EnumRotationStretchTo" : classes["InputString"],
  549. "EnumTrackAxis" : classes["InputString"],
  550. "EnumUpAxis" : classes["InputString"],
  551. "EnumLockAxis" : classes["InputString"],
  552. "EnumLimitMode" : classes["InputString"],
  553. "EnumYScaleMode" : classes["InputString"],
  554. "EnumXZScaleMode" : classes["InputString"],
  555. "EnumCurveSocket" : classes["InputString"],
  556. "EnumMetaRigSocket" : classes["InputString"],
  557. # Deformers
  558. "EnumSkinning" : classes["InputString"],
  559. #
  560. "FloatSocket" : classes["InputFloat"],
  561. "FloatFactorSocket" : classes["InputFloat"],
  562. "FloatPositiveSocket" : classes["InputFloat"],
  563. "FloatAngleSocket" : classes["InputFloat"],
  564. "VectorSocket" : classes["InputVector"],
  565. "VectorEulerSocket" : classes["InputVector"],
  566. "VectorTranslationSocket" : classes["InputVector"],
  567. "VectorScaleSocket" : classes["InputVector"],
  568. # Drivers
  569. "EnumDriverVariableType" : classes["InputString"],
  570. "EnumDriverVariableEvaluationSpace" : classes["InputString"],
  571. "EnumDriverRotationMode" : classes["InputString"],
  572. "EnumDriverType" : classes["InputString"],
  573. "EnumKeyframeInterpTypeSocket" : classes["InputString"],
  574. "EnumKeyframeBezierHandleTypeSocket" : classes["InputString"],
  575. # Math
  576. "MathFloatOperation" : classes["InputString"],
  577. "MathVectorOperation" : classes["InputString"],
  578. "MatrixTransformOperation" : classes["InputString"],
  579. # Schema
  580. "WildcardSocket" : None,
  581. }
  582. return socket_class_map.get(socket.bl_idname, None)
  583. ####################################
  584. # CURVE STUFF
  585. ####################################
  586. def make_perpendicular(v1, v2):
  587. projected = (v2.dot(v1) / v1.dot(v1)) * v1
  588. perpendicular = v2 - projected
  589. return perpendicular
  590. # this stuff could be branchless but I don't use it much TODO
  591. def cap(val, maxValue):
  592. if (val > maxValue):
  593. return maxValue
  594. return val
  595. def capMin(val, minValue):
  596. if (val < minValue):
  597. return minValue
  598. return val
  599. def wrap(min : float, max : float, value: float) -> float:
  600. range = max-min; remainder = value % range
  601. if remainder > max: return min + remainder-max
  602. else: return remainder
  603. def lerpVal(a, b, fac = 0.5):
  604. return a + ( (b-a) * fac)
  605. #wtf this doesn't do anything even remotely similar to wrap
  606. # HACK BAD FIXME UNBREAK ME BAD
  607. # I don't understand what this function does but I am using it in multiple places?
  608. def old_bad_wrap_that_should_be_refactored(val, maxValue, minValue = None):
  609. if (val > maxValue):
  610. return (-1 * ((maxValue - val) + 1))
  611. if ((minValue) and (val < minValue)):
  612. return (val + maxValue)
  613. return val
  614. #TODO clean this up
  615. def RibbonMeshEdgeLengths(m, ribbon):
  616. tE = ribbon[0]; bE = ribbon[1]; c = ribbon[2]
  617. lengths = []
  618. for i in range( len( tE ) ): #tE and bE are same length
  619. if (c == True):
  620. v1NextInd = tE[old_bad_wrap_that_should_be_refactored((i+1), len(tE) - 1)]
  621. else:
  622. v1NextInd = tE[cap((i+1) , len(tE) - 1 )]
  623. v1 = m.vertices[tE[i]]; v1Next = m.vertices[v1NextInd]
  624. if (c == True):
  625. v2NextInd = bE[old_bad_wrap_that_should_be_refactored((i+1), len(bE) - 1)]
  626. else:
  627. v2NextInd = bE[cap((i+1) , len(bE) - 1 )]
  628. v2 = m.vertices[bE[i]]; v2Next = m.vertices[v2NextInd]
  629. v = v1.co.lerp(v2.co, 0.5); vNext = v1Next.co.lerp(v2Next.co, 0.5)
  630. # get the center, edges may not be straight so total length
  631. # of one edge may be more than the ribbon center's length
  632. lengths.append(( v - vNext ).length)
  633. return lengths
  634. def EnsureCurveIsRibbon(crv, defaultRadius = 0.1):
  635. crvRadius = 0
  636. crv.data.offset = 0
  637. if (crv.data.bevel_depth == 0):
  638. crvRadius = crv.data.extrude
  639. else: #Set ribbon from bevel depth
  640. crvRadius = crv.data.bevel_depth
  641. crv.data.bevel_depth = 0
  642. crv.data.extrude = crvRadius
  643. if (crvRadius == 0):
  644. crv.data.extrude = defaultRadius
  645. def SetRibbonData(m, ribbon):
  646. #maybe this could be incorporated into the DetectWireEdges function?
  647. #maybe I can check for closed poly curves here? under what other circumstance
  648. # will I find the ends of the wire have identical coordinates?
  649. ribbonData = []
  650. tE = ribbon[0].copy(); bE = ribbon[1].copy()# circle = ribbon[2]
  651. #
  652. lengths = RibbonMeshEdgeLengths(m, ribbon)
  653. lengths.append(0)
  654. totalLength = sum(lengths)
  655. # m.calc_normals() #calculate normals
  656. # it appears this has been removed.
  657. for i, (t, b) in enumerate(zip(tE, bE)):
  658. ind = old_bad_wrap_that_should_be_refactored( (i + 1), len(tE) - 1 )
  659. tNext = tE[ind]; bNext = bE[ind]
  660. ribbonData.append( ( (t,b), (tNext, bNext), lengths[i] ) )
  661. #if this is a circle, the last v in vertData has a length, otherwise 0
  662. return ribbonData, totalLength
  663. def WireMeshEdgeLengths(m, wire):
  664. circle = False
  665. vIndex = wire.copy()
  666. for e in m.edges:
  667. if ((e.vertices[0] == vIndex[-1]) and (e.vertices[1] == vIndex[0])):
  668. #this checks for an edge between the first and last vertex in the wire
  669. circle = True
  670. break
  671. lengths = []
  672. for i in range(len(vIndex)):
  673. v = m.vertices[vIndex[i]]
  674. if (circle == True):
  675. vNextInd = vIndex[old_bad_wrap_that_should_be_refactored((i+1), len(vIndex) - 1)]
  676. else:
  677. vNextInd = vIndex[cap((i+1), len(vIndex) - 1 )]
  678. vNext = m.vertices[vNextInd]
  679. lengths.append(( v.co - vNext.co ).length)
  680. #if this is a circular wire mesh, this should wrap instead of cap
  681. return lengths
  682. def GetDataFromWire(m, wire):
  683. vertData = []
  684. vIndex = wire.copy()
  685. lengths = WireMeshEdgeLengths(m, wire)
  686. lengths.append(0)
  687. totalLength = sum(lengths)
  688. for i, vInd in enumerate(vIndex):
  689. #-1 to avoid IndexError
  690. vNext = vIndex[ (old_bad_wrap_that_should_be_refactored(i+1, len(vIndex) - 1)) ]
  691. vertData.append((vInd, vNext, lengths[i]))
  692. #if this is a circle, the last v in vertData has a length, otherwise 0
  693. return vertData, totalLength
  694. def DetectWireEdges(mesh):
  695. # Returns a list of vertex indices belonging to wire meshes
  696. # NOTE: this assumes a mesh object with only wire meshes
  697. ret = []
  698. import bmesh
  699. bm = bmesh.new()
  700. try:
  701. bm.from_mesh(mesh)
  702. ends = []
  703. for v in bm.verts:
  704. if (len(v.link_edges) == 1):
  705. ends.append(v.index)
  706. for e in bm.edges:
  707. assert (e.is_wire == True),"This function can only run on wire meshes"
  708. if (e.verts[1].index - e.verts[0].index != 1):
  709. ends.append(e.verts[1].index)
  710. ends.append(e.verts[0].index)
  711. for i in range(len(ends)//2): # // is floor division
  712. beg = ends[i*2]
  713. end = ends[(i*2)+1]
  714. indices = [(j + beg) for j in range ((end - beg) + 1)]
  715. ret.append(indices)
  716. finally:
  717. bm.free()
  718. return ret
  719. def FindNearestPointOnWireMesh(m, pointsList):
  720. from mathutils import Vector
  721. from mathutils.geometry import intersect_point_line
  722. from math import sqrt
  723. wires = DetectWireEdges(m)
  724. ret = []
  725. # prevFactor = None
  726. for wire, points in zip(wires, pointsList):
  727. vertData, total_length = GetDataFromWire(m, wire)
  728. factorsOut = []
  729. for p in points:
  730. prevDist = float('inf')
  731. curDist = float('inf')
  732. v1 = None
  733. v2 = None
  734. for i in range(len(vertData) - 1):
  735. #but it shouldn't check the last one
  736. if (p == m.vertices[i].co):
  737. v1 = vertData[i]
  738. v2 = vertData[i+1]
  739. offset = 0
  740. break
  741. else:
  742. curDist = ( ((m.vertices[vertData[i][0]].co - p).length) +
  743. ((m.vertices[vertData[i][1]].co - p).length) )/2
  744. if (curDist < prevDist):
  745. v1 = vertData[i]
  746. v2 = vertData[i+1]
  747. prevDist = curDist
  748. offset = intersect_point_line(p, m.vertices[v1[0]].co,
  749. m.vertices[v2[0]].co)[1]
  750. if (offset < 0):
  751. offset = 0
  752. elif (offset > 1):
  753. offset = 1
  754. # Assume the vertices are in order
  755. v1Length = 0
  756. v2Length = v2[2]
  757. for i in range(v1[0]):
  758. v1Length += vertData[i][2]
  759. factor = ((offset * (v2Length)) + v1Length )/total_length
  760. factor = wrap(0, 1, factor) # doesn't hurt to wrap it if it's over 1 or less than 0
  761. factorsOut.append(factor)
  762. ret.append( factorsOut )
  763. return ret
  764. def mesh_from_curve(crv, context, ribbon=True):
  765. """Utility function for converting a mesh to a curve
  766. which will return the correct mesh even with modifiers"""
  767. import bpy
  768. m = None
  769. bevel = crv.data.bevel_depth
  770. extrude = crv.data.extrude
  771. offset = crv.data.offset
  772. try:
  773. if (len(crv.modifiers) > 0):
  774. do_unlink = False
  775. if (not context.scene.collection.all_objects.get(crv.name)):
  776. context.collection.objects.link(crv) # i guess this forces the dg to update it?
  777. do_unlink = True
  778. dg = context.view_layer.depsgraph
  779. # just gonna modify it for now lol
  780. if ribbon:
  781. EnsureCurveIsRibbon(crv)
  782. else:
  783. crv.data.bevel_depth=0
  784. crv.data.extrude=0
  785. crv.data.offset=0
  786. # try:
  787. dg.update()
  788. mOb = crv.evaluated_get(dg)
  789. m = bpy.data.meshes.new_from_object(mOb)
  790. m.name=crv.data.name+'_mesh'
  791. if (do_unlink):
  792. context.collection.objects.unlink(crv)
  793. else: # (ಥ﹏ಥ) why can't I just use this !
  794. # for now I will just do it like this
  795. if ribbon:
  796. EnsureCurveIsRibbon(crv)
  797. else:
  798. crv.data.bevel_depth=0
  799. crv.data.extrude=0
  800. crv.data.offset=0
  801. m = bpy.data.meshes.new_from_object(crv)
  802. finally:
  803. crv.data.bevel_depth = bevel
  804. crv.data.extrude = extrude
  805. crv.data.offset = offset
  806. return m
  807. def DetectRibbon(f, bm, skipMe):
  808. fFirst = f.index
  809. cont = True
  810. circle = False
  811. tEdge, bEdge = [],[]
  812. while (cont == True):
  813. skipMe.add(f.index)
  814. tEdge.append (f.loops[0].vert.index) # top-left
  815. bEdge.append (f.loops[3].vert.index) # bottom-left
  816. nEdge = bm.edges.get([f.loops[1].vert, f.loops[2].vert])
  817. nFaces = nEdge.link_faces
  818. if (len(nFaces) == 1):
  819. cont = False
  820. else:
  821. for nFace in nFaces:
  822. if (nFace != f):
  823. f = nFace
  824. break
  825. if (f.index == fFirst):
  826. cont = False
  827. circle = True
  828. if (cont == False): # we've reached the end, get the last two:
  829. tEdge.append (f.loops[1].vert.index) # top-right
  830. bEdge.append (f.loops[2].vert.index) # bottom-right
  831. # this will create a loop for rings --
  832. # "the first shall be the last and the last shall be first"
  833. return (tEdge,bEdge,circle)
  834. def DetectRibbons(m, fReport = None):
  835. # Returns list of vertex indices belonging to ribbon mesh edges
  836. # NOTE: this assumes a mesh object with only ribbon meshes
  837. # ---DO NOT call this script with a mesh that isn't a ribbon!--- #
  838. import bmesh
  839. bm = bmesh.new()
  840. bm.from_mesh(m)
  841. mIslands, mIsland = [], []
  842. skipMe = set()
  843. bm.faces.ensure_lookup_table()
  844. #first, get a list of mesh islands
  845. for f in bm.faces:
  846. if (f.index in skipMe):
  847. continue #already done here
  848. checkMe = [f]
  849. while (len(checkMe) > 0):
  850. facesFound = 0
  851. for f in checkMe:
  852. if (f.index in skipMe):
  853. continue #already done here
  854. mIsland.append(f)
  855. skipMe.add(f.index)
  856. for e in f.edges:
  857. checkMe += e.link_faces
  858. if (facesFound == 0):
  859. #this is the last iteration
  860. mIslands.append(mIsland)
  861. checkMe, mIsland = [], []
  862. ribbons = []
  863. skipMe = set() # to store ends already checked
  864. for mIsl in mIslands:
  865. ribbon = None
  866. first = float('inf')
  867. for f in mIsl:
  868. if (f.index in skipMe):
  869. continue #already done here
  870. if (f.index < first):
  871. first = f.index
  872. adjF = 0
  873. for e in f.edges:
  874. adjF+= (len(e.link_faces) - 1)
  875. # every face other than this one is added to the list
  876. if (adjF == 1):
  877. ribbon = (DetectRibbon(f, bm, skipMe) )
  878. break
  879. if (ribbon == None):
  880. ribbon = (DetectRibbon(bm.faces[first], bm, skipMe) )
  881. ribbons.append(ribbon)
  882. # print (ribbons)
  883. return ribbons
  884. def data_from_ribbon_mesh(m, factorsList, mat, ribbons = None, fReport = None):
  885. #Note, factors list should be equal in length the the number of wires
  886. #Now working for multiple wires, ugly tho
  887. if (ribbons == None):
  888. ribbons = DetectRibbons(m, fReport=fReport)
  889. if (ribbons is None):
  890. if (fReport):
  891. fReport(type = {'ERROR'}, message="No ribbon to get data from.")
  892. else:
  893. print ("No ribbon to get data from.")
  894. return None
  895. ret = []
  896. for factors, ribbon in zip(factorsList, ribbons):
  897. points = []
  898. widths = []
  899. normals = []
  900. ribbonData, totalLength = SetRibbonData(m, ribbon)
  901. for fac in factors:
  902. if (fac == 0):
  903. data = ribbonData[0]
  904. curFac = 0
  905. elif (fac == 1):
  906. data = ribbonData[-1]
  907. curFac = 0
  908. else:
  909. targetLength = totalLength * fac
  910. data = ribbonData[0]
  911. curLength = 0
  912. for ( (t, b), (tNext, bNext), length,) in ribbonData:
  913. if (curLength >= targetLength):
  914. break
  915. curLength += length
  916. data = ( (t, b), (tNext, bNext), length,)
  917. targetLengthAtEdge = (curLength - targetLength)
  918. if (targetLength == 0):
  919. curFac = 0
  920. elif (targetLength == totalLength):
  921. curFac = 1
  922. else:
  923. # NOTE: This can be Zero. That should throw an error.
  924. curFac = 1 - (targetLengthAtEdge/ data[2]) #length
  925. t1 = m.vertices[data[0][0]]; b1 = m.vertices[data[0][1]]
  926. t2 = m.vertices[data[1][0]]; b2 = m.vertices[data[1][1]]
  927. #location
  928. loc1 = (t1.co).lerp(b1.co, 0.5)
  929. loc2 = (t2.co).lerp(b2.co, 0.5)
  930. #width
  931. w1 = (t1.co - b1.co).length/2
  932. w2 = (t2.co - b2.co).length/2 #radius, not diameter
  933. #normal
  934. n1 = (t1.normal).slerp(b1.normal, 0.5)
  935. n2 = (t1.normal).slerp(b2.normal, 0.5)
  936. if ((data[0][0] > data[1][0]) and (ribbon[2] == False)):
  937. curFac = 0
  938. #don't interpolate if at the end of a ribbon that isn't circular
  939. if ( 0 < curFac < 1):
  940. outPoint = loc1.lerp(loc2, curFac)
  941. outNorm = n1.lerp(n2, curFac)
  942. outWidth = w1 + ( (w2-w1) * curFac)
  943. elif (curFac <= 0):
  944. outPoint = loc1.copy()
  945. outNorm = n1
  946. outWidth = w1
  947. elif (curFac >= 1):
  948. outPoint = loc2.copy()
  949. outNorm = n2
  950. outWidth = w2
  951. outPoint = mat @ outPoint
  952. outNorm.normalize()
  953. points.append ( outPoint.copy() ) #copy because this is an actual vertex location
  954. widths.append ( outWidth )
  955. normals.append( outNorm )
  956. ret.append( (points, widths, normals) )
  957. return ret # this is a list of tuples containing three lists
  958. #This bisection search is generic, and it searches based on the
  959. # magnitude of the error, rather than the sign.
  960. # If the sign of the error is meaningful, a simpler function
  961. # can be used.
  962. def do_bisect_search_by_magnitude(
  963. owner,
  964. attribute,
  965. index = None,
  966. test_function = None,
  967. modify = None,
  968. max_iterations = 10000,
  969. threshold = 0.0001,
  970. thresh2 = 0.0005,
  971. context = None,
  972. update_dg = None,
  973. ):
  974. from math import floor
  975. i = 0; best_so_far = 0; best = float('inf')
  976. min = 0; center = max_iterations//2; max = max_iterations
  977. # enforce getting the absolute value, in case the function has sign information
  978. # The sign may be useful in a sign-aware bisect search, but this one is more robust!
  979. test = lambda : abs(test_function(owner, attribute, index, context = context,))
  980. while (i <= max_iterations):
  981. upper = (max - ((max-center))//2)
  982. modify(owner, attribute, index, upper, context = context); error1 = test()
  983. lower = (center - ((center-min))//2)
  984. modify(owner, attribute, index, lower, context = context); error2 = test()
  985. if (error1 < error2):
  986. min = center
  987. center, check = upper, upper
  988. error = error1
  989. else:
  990. max = center
  991. center, check = lower, lower
  992. error = error2
  993. if (error <= threshold) or (min == max-1):
  994. break
  995. if (error < thresh2):
  996. j = min
  997. while (j < max):
  998. modify(owner, attribute, index, j * 1/max_iterations, context = context)
  999. error = test()
  1000. if (error < best):
  1001. best_so_far = j; best = error
  1002. if (error <= threshold):
  1003. break
  1004. j+=1
  1005. else: # loop has completed without finding a solution
  1006. i = best_so_far; error = test()
  1007. modify(owner, attribute, index, best_so_far, context = context)
  1008. break
  1009. if (error < best):
  1010. best_so_far = check; best = error
  1011. i+=1
  1012. if update_dg:
  1013. update_dg.update()
  1014. else: # Loop has completed without finding a solution
  1015. i = best_so_far
  1016. modify(owner, attribute, best_so_far, context = context); i+=1