utilities.py 39 KB

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