rna data path names which are more likely to break animations.
[blender-staging.git] / release / scripts / modules / rigify_utils.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 # rigify its self does not depend on this module, however some of the
22 # rigify templates use these utility functions.
23 #
24 # So even though this can be for general purpose use, this module was created
25 # for rigify so in some cases seemingly generic functions make assumptions
26 # that a generic function would need to check for.
27
28 import bpy
29 from mathutils import Vector
30 from rna_prop_ui import rna_idprop_ui_prop_get
31
32 DELIMITER = '-._'
33 EMPTY_LAYER = [False] * 32
34
35
36 def add_stretch_to(obj, from_name, to_name, name):
37     '''
38     Adds a bone that stretches from one to another
39     '''
40
41     mode_orig = obj.mode
42     bpy.ops.object.mode_set(mode='EDIT')
43
44     arm = obj.data
45     stretch_ebone = arm.edit_bones.new(name)
46     stretch_name = stretch_ebone.name
47     del name
48
49     head = stretch_ebone.head = arm.edit_bones[from_name].head.copy()
50     #tail = stretch_ebone.tail = arm.edit_bones[to_name].head.copy()
51
52     # annoying exception for zero length bones, since its using stretch_to the rest pose doesnt really matter
53     #if (head - tail).length < 0.1:
54     if 1:
55         tail = stretch_ebone.tail = arm.edit_bones[from_name].tail.copy()
56
57
58     # Now for the constraint
59     bpy.ops.object.mode_set(mode='OBJECT')
60
61     stretch_pbone = obj.pose.bones[stretch_name]
62
63     con = stretch_pbone.constraints.new('COPY_LOCATION')
64     con.target = obj
65     con.subtarget = from_name
66
67     con = stretch_pbone.constraints.new('STRETCH_TO')
68     con.target = obj
69     con.subtarget = to_name
70     con.rest_length = (head - tail).length
71     con.keep_axis = 'PLANE_X'
72     con.volume = 'NO_VOLUME'
73
74     bpy.ops.object.mode_set(mode=mode_orig)
75
76     return stretch_name
77
78
79 def copy_bone_simple(arm, from_bone, name, parent=False):
80     ebone = arm.edit_bones[from_bone]
81     ebone_new = arm.edit_bones.new(name)
82
83     if parent:
84         ebone_new.use_connect = ebone.use_connect
85         ebone_new.parent = ebone.parent
86
87     ebone_new.head = ebone.head
88     ebone_new.tail = ebone.tail
89     ebone_new.roll = ebone.roll
90     ebone_new.layers = list(ebone.layers)
91     return ebone_new
92
93
94 def copy_bone_simple_list(arm, from_bones, to_bones, parent=False):
95
96     if len(from_bones) != len(to_bones):
97         raise Exception("bone list sizes must match")
98
99     copy_bones = [copy_bone_simple(arm, bone_name, to_bones[i], True) for i, bone_name in enumerate(from_bones)]
100
101     # now we need to re-parent
102     for ebone in copy_bones:
103         parent = ebone.parent
104         if parent:
105             try:
106                 i = from_bones.index(parent.name)
107             except:
108                 i = -1
109
110             if i == -1:
111                 ebone.parent = None
112             else:
113                 ebone.parent = copy_bones[i]
114
115     return copy_bones
116
117
118 def blend_bone_list(obj, apply_bones, from_bones, to_bones, target_bone=None, target_prop="blend", blend_default=0.5):
119
120     if obj.mode == 'EDIT':
121         raise Exception("blending cant be called in editmode")
122
123     if len(apply_bones) != len(from_bones):
124         raise Exception("lists differ in length (from -> apply): \n\t%s\n\t%s" % (from_bones, apply_bones))
125     if len(apply_bones) != len(to_bones):
126         raise Exception("lists differ in length (to -> apply): \n\t%s\n\t%s" % (to_bones, apply_bones))
127
128     # setup the blend property
129     if target_bone is None:
130         target_bone = apply_bones[-1] # default to the last bone
131
132     prop_pbone = obj.pose.bones[target_bone]
133     if prop_pbone.get(target_bone) is None:
134         prop = rna_idprop_ui_prop_get(prop_pbone, target_prop, create=True)
135         prop_pbone[target_prop] = blend_default
136         prop["soft_min"] = 0.0
137         prop["soft_max"] = 1.0
138
139     driver_path = prop_pbone.path_from_id() + ('["%s"]' % target_prop)
140
141     def blend_target(driver):
142         var = driver.variables.new()
143         var.name = target_bone
144         var.targets[0].id_type = 'OBJECT'
145         var.targets[0].id = obj
146         var.targets[0].data_path = driver_path
147
148     def blend_transforms(new_pbone, from_bone_name, to_bone_name):
149         con = new_pbone.constraints.new('COPY_TRANSFORMS')
150         con.target = obj
151         con.subtarget = from_bone_name
152
153         con = new_pbone.constraints.new('COPY_TRANSFORMS')
154         con.target = obj
155         con.subtarget = to_bone_name
156
157         fcurve = con.driver_add("influence")
158         driver = fcurve.driver
159         driver.type = 'AVERAGE'
160         fcurve.modifiers.remove(0) # grr dont need a modifier
161
162         blend_target(driver)
163
164     for i, new_bone_name in enumerate(apply_bones):
165         from_bone_name = from_bones[i]
166         to_bone_name = to_bones[i]
167
168         # allow skipping some bones by having None in the list
169         if None in (new_bone_name, from_bone_name, to_bone_name):
170             continue
171
172         new_pbone = obj.pose.bones[new_bone_name]
173
174         blend_transforms(new_pbone, from_bone_name, to_bone_name)
175
176
177 def add_pole_target_bone(obj, base_bone_name, name, mode='CROSS'):
178     '''
179     Does not actually create a poll target, just the bone to use as a poll target
180     '''
181     mode_orig = obj.mode
182     bpy.ops.object.mode_set(mode='EDIT')
183
184     arm = obj.data
185
186     poll_ebone = arm.edit_bones.new(name)
187     base_ebone = arm.edit_bones[base_bone_name]
188     poll_name = poll_ebone.name
189     parent_ebone = base_ebone.parent
190
191     base_head = base_ebone.head.copy()
192     base_tail = base_ebone.tail.copy()
193     base_dir = base_head - base_tail
194
195     parent_head = parent_ebone.head.copy()
196     parent_tail = parent_ebone.tail.copy()
197     parent_dir = parent_head - parent_tail
198
199     distance = (base_dir.length + parent_dir.length)
200
201     if mode == 'CROSS':
202         # direction from the angle of the joint
203         offset = base_dir.copy().normalize() - parent_dir.copy().normalize()
204         offset.length = distance
205     elif mode == 'ZAVERAGE':
206         # between both bones Z axis
207         z_axis_a = base_ebone.matrix.copy().rotation_part() * Vector((0.0, 0.0, -1.0))
208         z_axis_b = parent_ebone.matrix.copy().rotation_part() * Vector((0.0, 0.0, -1.0))
209         offset = (z_axis_a + z_axis_b).normalize() * distance
210     else:
211         # preset axis
212         offset = Vector((0.0, 0.0, 0.0))
213         if mode[0] == "+":
214             val = distance
215         else:
216             val = - distance
217
218         setattr(offset, mode[1].lower(), val)
219
220     poll_ebone.head = base_head + offset
221     poll_ebone.tail = base_head + (offset * (1.0 - (1.0 / 4.0)))
222
223     bpy.ops.object.mode_set(mode=mode_orig)
224
225     return poll_name
226
227
228 def get_side_name(name):
229     '''
230     Returns the last part of a string (typically a bone's name) indicating
231     whether it is a a left or right (or center, or whatever) bone.
232     Returns an empty string if nothing is found.
233     '''
234     if name[-2] in DELIMITER:
235         return name[-2:]
236     else:
237         return ""
238
239
240 def get_base_name(name):
241     '''
242     Returns the part of a string (typically a bone's name) corresponding to it's
243     base name (no sidedness, no ORG prefix).
244     '''
245     if name[-2] in DELIMITER:
246         return name[:-2]
247     else:
248         return name
249
250
251 def write_meta_rig(obj, func_name="metarig_template"):
252     '''
253     Write a metarig as a python script, this rig is to have all info needed for
254     generating the real rig with rigify.
255     '''
256     code = []
257
258     code.append("def %s():" % func_name)
259     code.append("    # generated by rigify.write_meta_rig")
260     bpy.ops.object.mode_set(mode='EDIT')
261     code.append("    bpy.ops.object.mode_set(mode='EDIT')")
262
263     code.append("    obj = bpy.context.active_object")
264     code.append("    arm = obj.data")
265
266     arm = obj.data
267     # write parents first
268     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
269     bones.sort(key=lambda item: item[0])
270     bones = [item[1] for item in bones]
271
272
273     for bone_name in bones:
274         bone = arm.edit_bones[bone_name]
275         code.append("    bone = arm.edit_bones.new('%s')" % bone.name)
276         code.append("    bone.head[:] = %.4f, %.4f, %.4f" % bone.head.to_tuple(4))
277         code.append("    bone.tail[:] = %.4f, %.4f, %.4f" % bone.tail.to_tuple(4))
278         code.append("    bone.roll = %.4f" % bone.roll)
279         code.append("    bone.use_connect = %s" % str(bone.use_connect))
280         if bone.parent:
281             code.append("    bone.parent = arm.edit_bones['%s']" % bone.parent.name)
282
283     bpy.ops.object.mode_set(mode='OBJECT')
284     code.append("")
285     code.append("    bpy.ops.object.mode_set(mode='OBJECT')")
286
287     for bone_name in bones:
288         pbone = obj.pose.bones[bone_name]
289         pbone_written = False
290
291         # Only 1 level of props, simple types supported
292         for key, value in pbone.items():
293             if key.startswith("_"):
294                 continue
295
296             if type(value) not in (float, str, int):
297                 print("Unsupported ID Prop:", str((key, value)))
298                 continue
299
300             if type(value) == str:
301                 value = "'" + value + "'"
302
303             if not pbone_written: # only write bones we need
304                 code.append("    pbone = obj.pose.bones['%s']" % bone_name)
305
306             code.append("    pbone['%s'] = %s" % (key, value))
307
308     return "\n".join(code)
309
310
311 # *** bone class collection ***
312
313
314 def bone_class_instance(obj, slots, name="BoneContainer"):
315     '''
316     bone collection utility class to help manage cases with
317     edit/pose/bone bones where switching modes can invalidate some of the members.
318
319     there are also utility functions for manipulating all members.
320     '''
321
322     attr_names = tuple(slots) # dont modify the original
323
324     if len(slots) != len(set(slots)):
325         raise Exception("duplicate entries found %s" % attr_names)
326
327     slots = list(slots) # dont modify the original
328     for i in range(len(slots)):
329         member = slots[i]
330         slots.append(member + "_b") # bone bone
331         slots.append(member + "_p") # pose bone
332         slots.append(member + "_e") # edit bone
333
334     class_dict = { \
335         "obj": obj, \
336         "attr_names": attr_names, \
337         "attr_initialize": _bone_class_instance_attr_initialize, \
338         "update": _bone_class_instance_update, \
339         "rename": _bone_class_instance_rename, \
340         "names": _bone_class_instance_names, \
341         "copy": _bone_class_instance_copy, \
342         "blend": _bone_class_instance_blend, \
343     }
344
345     instance = auto_class_instance(slots, name, class_dict)
346     return instance
347
348
349 def auto_class(slots, name="ContainerClass", class_dict=None):
350
351     if class_dict:
352         class_dict = class_dict.copy()
353     else:
354         class_dict = {}
355
356     class_dict["__slots__"] = tuple(slots)
357
358     return type(name, (object,), class_dict)
359
360
361 def auto_class_instance(slots, name="ContainerClass", class_dict=None):
362     return auto_class(slots, name, class_dict)()
363
364
365 def _bone_class_instance_attr_initialize(self, attr_names, bone_names):
366     ''' Initializes attributes, both lists must be aligned
367     '''
368     for attr in self.attr_names:
369         i = attr_names.index(attr)
370         setattr(self, attr, bone_names[i])
371
372     self.update()
373
374
375 def _bone_class_instance_update(self):
376     ''' Re-Assigns bones from the blender data
377     '''
378     arm = self.obj.data
379     bbones = arm.bones
380     pbones = self.obj.pose.bones
381     ebones = arm.edit_bones
382
383     for member in self.attr_names:
384         name = getattr(self, member, None)
385         if name is not None:
386             setattr(self, member + "_b", bbones.get(name))
387             setattr(self, member + "_p", pbones.get(name))
388             setattr(self, member + "_e", ebones.get(name))
389
390
391 def _bone_class_instance_rename(self, attr, new_name):
392     ''' Rename bones, editmode only
393     '''
394
395     if self.obj.mode != 'EDIT':
396         raise Exception("Only rename in editmode supported")
397
398     ebone = getattr(self, attr + "_e")
399     ebone.name = new_name
400
401     # we may not get what is asked for so get the name from the editbone
402     setattr(self, attr, ebone.name)
403
404
405 def _bone_class_instance_copy(self, from_fmt="%s", to_fmt="%s", exclude_attrs=(), base_names=None):
406     from_name_ls = []
407     new_name_ls = []
408     new_slot_ls = []
409
410     for attr in self.attr_names:
411
412         if attr in exclude_attrs:
413             continue
414
415         bone_name_orig = getattr(self, attr)
416         ebone = getattr(self, attr + "_e")
417         # orig_names[attr] = bone_name_orig
418
419         # insert formatting
420         if from_fmt != "%s":
421             bone_name = from_fmt % bone_name_orig
422             ebone.name = bone_name
423             bone_name = ebone.name # cant be sure we get what we ask for
424         else:
425             bone_name = bone_name_orig
426
427         setattr(self, attr, bone_name)
428
429         new_slot_ls.append(attr)
430         from_name_ls.append(bone_name)
431         if base_names:
432             bone_name_orig = base_names[bone_name_orig]
433         new_name_ls.append(to_fmt % bone_name_orig)
434
435     new_bones = copy_bone_simple_list(self.obj.data, from_name_ls, new_name_ls, True)
436     new_bc = bone_class_instance(self.obj, new_slot_ls)
437
438     for i, attr in enumerate(new_slot_ls):
439         ebone = new_bones[i]
440         setattr(new_bc, attr + "_e", ebone)
441         setattr(new_bc, attr, ebone.name)
442
443     return new_bc
444
445
446 def _bone_class_instance_names(self):
447     return [getattr(self, attr) for attr in self.attr_names]
448
449
450 def _bone_class_instance_blend(self, from_bc, to_bc, target_bone=None, target_prop="blend"):
451     '''
452     Use for blending bone chains.
453
454     blend_target = (bone_name, bone_property)
455     default to the last bone, blend prop
456
457     XXX - toggles editmode, need to re-validate all editbones :(
458     '''
459
460     if self.attr_names != from_bc.attr_names or self.attr_names != to_bc.attr_names:
461         raise Exception("can only blend between matching chains")
462
463     apply_bones = [getattr(self, attr) for attr in self.attr_names]
464     from_bones = [getattr(from_bc, attr) for attr in from_bc.attr_names]
465     to_bones = [getattr(to_bc, attr) for attr in to_bc.attr_names]
466
467     blend_bone_list(self.obj, apply_bones, from_bones, to_bones, target_bone, target_prop)