rigify/rig_ui_template.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 # <pep8 compliant>
   4
   5 import bpy
   6
   7 from collections import OrderedDict
   8
   9 from .utils.animation import SCRIPT_REGISTER_BAKE, SCRIPT_UTILITIES_BAKE
  10
  11 from . import base_generate
  12
  13 from rna_prop_ui import rna_idprop_quote_path
  14
  15
  16 UI_IMPORTS = [
  17     'import bpy',
  18     'import math',
  19     'import json',
  20     'import collections',
  21     'import traceback',
  22     'from math import pi',
  23     'from bpy.props import StringProperty',
  24     'from mathutils import Euler, Matrix, Quaternion, Vector',
  25     'from rna_prop_ui import rna_idprop_quote_path',
  26 ]
  27
  28 UI_BASE_UTILITIES = '''
  29 rig_id = "%s"
  30
  31
  32 ############################
  33 ## Math utility functions ##
  34 ############################
  35
  36 def perpendicular_vector(v):
  37     """ Returns a vector that is perpendicular to the one given.
  38         The returned vector is _not_ guaranteed to be normalized.
  39     """
  40     # Create a vector that is not aligned with v.
  41     # It doesn't matter what vector.  Just any vector
  42     # that's guaranteed to not be pointing in the same
  43     # direction.
  44     if abs(v[0]) < abs(v[1]):
  45         tv = Vector((1,0,0))
  46     else:
  47         tv = Vector((0,1,0))
  48
  49     # Use cross prouct to generate a vector perpendicular to
  50     # both tv and (more importantly) v.
  51     return v.cross(tv)
  52
  53
  54 def rotation_difference(mat1, mat2):
  55     """ Returns the shortest-path rotational difference between two
  56         matrices.
  57     """
  58     q1 = mat1.to_quaternion()
  59     q2 = mat2.to_quaternion()
  60     angle = math.acos(min(1,max(-1,q1.dot(q2)))) * 2
  61     if angle > pi:
  62         angle = -angle + (2*pi)
  63     return angle
  64
  65 def find_min_range(f,start_angle,delta=pi/8):
  66     """ finds the range where lies the minimum of function f applied on bone_ik and bone_fk
  67         at a certain angle.
  68     """
  69     angle = start_angle
  70     while (angle > (start_angle - 2*pi)) and (angle < (start_angle + 2*pi)):
  71         l_dist = f(angle-delta)
  72         c_dist = f(angle)
  73         r_dist = f(angle+delta)
  74         if min((l_dist,c_dist,r_dist)) == c_dist:
  75             return (angle-delta,angle+delta)
  76         else:
  77             angle=angle+delta
  78
  79 def ternarySearch(f, left, right, absolutePrecision):
  80     """
  81     Find minimum of unimodal function f() within [left, right]
  82     To find the maximum, revert the if/else statement or revert the comparison.
  83     """
  84     while True:
  85         #left and right are the current bounds; the maximum is between them
  86         if abs(right - left) < absolutePrecision:
  87             return (left + right)/2
  88
  89         leftThird = left + (right - left)/3
  90         rightThird = right - (right - left)/3
  91
  92         if f(leftThird) > f(rightThird):
  93             left = leftThird
  94         else:
  95             right = rightThird
  96 '''
  97
  98 UTILITIES_FUNC_COMMON_IKFK = ['''
  99 #########################################
 100 ## "Visual Transform" helper functions ##
 101 #########################################
 102
 103 def get_pose_matrix_in_other_space(mat, pose_bone):
 104     """ Returns the transform matrix relative to pose_bone's current
 105         transform space.  In other words, presuming that mat is in
 106         armature space, slapping the returned matrix onto pose_bone
 107         should give it the armature-space transforms of mat.
 108     """
 109     return pose_bone.id_data.convert_space(matrix=mat, pose_bone=pose_bone, from_space='POSE', to_space='LOCAL')
 110
 111
 112 def convert_pose_matrix_via_rest_delta(mat, from_bone, to_bone):
 113     """Convert pose of one bone to another bone, preserving the rest pose difference between them."""
 114     return mat @ from_bone.bone.matrix_local.inverted() @ to_bone.bone.matrix_local
 115
 116
 117 def convert_pose_matrix_via_pose_delta(mat, from_bone, to_bone):
 118     """Convert pose of one bone to another bone, preserving the current pose difference between them."""
 119     return mat @ from_bone.matrix.inverted() @ to_bone.matrix
 120
 121
 122 def get_local_pose_matrix(pose_bone):
 123     """ Returns the local transform matrix of the given pose bone.
 124     """
 125     return get_pose_matrix_in_other_space(pose_bone.matrix, pose_bone)
 126
 127
 128 def set_pose_translation(pose_bone, mat):
 129     """ Sets the pose bone's translation to the same translation as the given matrix.
 130         Matrix should be given in bone's local space.
 131     """
 132     pose_bone.location = mat.to_translation()
 133
 134
 135 def set_pose_rotation(pose_bone, mat):
 136     """ Sets the pose bone's rotation to the same rotation as the given matrix.
 137         Matrix should be given in bone's local space.
 138     """
 139     q = mat.to_quaternion()
 140
 141     if pose_bone.rotation_mode == 'QUATERNION':
 142         pose_bone.rotation_quaternion = q
 143     elif pose_bone.rotation_mode == 'AXIS_ANGLE':
 144         pose_bone.rotation_axis_angle[0] = q.angle
 145         pose_bone.rotation_axis_angle[1] = q.axis[0]
 146         pose_bone.rotation_axis_angle[2] = q.axis[1]
 147         pose_bone.rotation_axis_angle[3] = q.axis[2]
 148     else:
 149         pose_bone.rotation_euler = q.to_euler(pose_bone.rotation_mode)
 150
 151
 152 def set_pose_scale(pose_bone, mat):
 153     """ Sets the pose bone's scale to the same scale as the given matrix.
 154         Matrix should be given in bone's local space.
 155     """
 156     pose_bone.scale = mat.to_scale()
 157
 158
 159 def match_pose_translation(pose_bone, target_bone):
 160     """ Matches pose_bone's visual translation to target_bone's visual
 161         translation.
 162         This function assumes you are in pose mode on the relevant armature.
 163     """
 164     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 165     set_pose_translation(pose_bone, mat)
 166
 167
 168 def match_pose_rotation(pose_bone, target_bone):
 169     """ Matches pose_bone's visual rotation to target_bone's visual
 170         rotation.
 171         This function assumes you are in pose mode on the relevant armature.
 172     """
 173     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 174     set_pose_rotation(pose_bone, mat)
 175
 176
 177 def match_pose_scale(pose_bone, target_bone):
 178     """ Matches pose_bone's visual scale to target_bone's visual
 179         scale.
 180         This function assumes you are in pose mode on the relevant armature.
 181     """
 182     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 183     set_pose_scale(pose_bone, mat)
 184
 185
 186 ##############################
 187 ## IK/FK snapping functions ##
 188 ##############################
 189
 190 def correct_rotation(view_layer, bone_ik, target_matrix, *, ctrl_ik=None):
 191     """ Corrects the ik rotation in ik2fk snapping functions
 192     """
 193
 194     axis = target_matrix.to_3x3().col[1].normalized()
 195     ctrl_ik = ctrl_ik or bone_ik
 196
 197     def distance(angle):
 198         # Rotate the bone and return the actual angle between bones
 199         ctrl_ik.rotation_euler[1] = angle
 200         view_layer.update()
 201
 202         return -(bone_ik.vector.normalized().dot(axis))
 203
 204     if ctrl_ik.rotation_mode in {'QUATERNION', 'AXIS_ANGLE'}:
 205         ctrl_ik.rotation_mode = 'ZXY'
 206
 207     start_angle = ctrl_ik.rotation_euler[1]
 208
 209     alfarange = find_min_range(distance, start_angle)
 210     alfamin = ternarySearch(distance, alfarange[0], alfarange[1], pi / 180)
 211
 212     ctrl_ik.rotation_euler[1] = alfamin
 213     view_layer.update()
 214
 215
 216 def correct_scale(view_layer, bone_ik, target_matrix, *, ctrl_ik=None):
 217     """ Correct the scale of the base IK bone. """
 218     input_scale = target_matrix.to_scale()
 219     ctrl_ik = ctrl_ik or bone_ik
 220
 221     for i in range(3):
 222         cur_scale = bone_ik.matrix.to_scale()
 223
 224         ctrl_ik.scale = [
 225             v * i / c for v, i, c in zip(bone_ik.scale, input_scale, cur_scale)
 226         ]
 227
 228         view_layer.update()
 229
 230         if all(abs((c - i)/i) < 0.01 for i, c in zip(input_scale, cur_scale)):
 231             break
 232
 233
 234 def match_pole_target(view_layer, ik_first, ik_last, pole, match_bone_matrix, length):
 235     """ Places an IK chain's pole target to match ik_first's
 236         transforms to match_bone.  All bones should be given as pose bones.
 237         You need to be in pose mode on the relevant armature object.
 238         ik_first: first bone in the IK chain
 239         ik_last:  last bone in the IK chain
 240         pole:  pole target bone for the IK chain
 241         match_bone:  bone to match ik_first to (probably first bone in a matching FK chain)
 242         length:  distance pole target should be placed from the chain center
 243     """
 244     a = ik_first.matrix.to_translation()
 245     b = ik_last.matrix.to_translation() + ik_last.vector
 246
 247     # Vector from the head of ik_first to the
 248     # tip of ik_last
 249     ikv = b - a
 250
 251     # Get a vector perpendicular to ikv
 252     pv = perpendicular_vector(ikv).normalized() * length
 253
 254     def set_pole(pvi):
 255         """ Set pole target's position based on a vector
 256             from the arm center line.
 257         """
 258         # Translate pvi into armature space
 259         ploc = a + (ikv/2) + pvi
 260
 261         # Set pole target to location
 262         mat = get_pose_matrix_in_other_space(Matrix.Translation(ploc), pole)
 263         set_pose_translation(pole, mat)
 264
 265         view_layer.update()
 266
 267     set_pole(pv)
 268
 269     # Get the rotation difference between ik_first and match_bone
 270     angle = rotation_difference(ik_first.matrix, match_bone_matrix)
 271
 272     # Try compensating for the rotation difference in both directions
 273     pv1 = Matrix.Rotation(angle, 4, ikv) @ pv
 274     set_pole(pv1)
 275     ang1 = rotation_difference(ik_first.matrix, match_bone_matrix)
 276
 277     pv2 = Matrix.Rotation(-angle, 4, ikv) @ pv
 278     set_pole(pv2)
 279     ang2 = rotation_difference(ik_first.matrix, match_bone_matrix)
 280
 281     # Do the one with the smaller angle
 282     if ang1 < ang2:
 283         set_pole(pv1)
 284
 285 ##########
 286 ## Misc ##
 287 ##########
 288
 289 def parse_bone_names(names_string):
 290     if names_string[0] == '[' and names_string[-1] == ']':
 291         return eval(names_string)
 292     else:
 293         return names_string
 294
 295 ''']
 296
 297 UTILITIES_FUNC_OLD_ARM_FKIK = ['''
 298 ######################
 299 ## IK Arm functions ##
 300 ######################
 301
 302 def fk2ik_arm(obj, fk, ik):
 303     """ Matches the fk bones in an arm rig to the ik bones.
 304         obj: armature object
 305         fk:  list of fk bone names
 306         ik:  list of ik bone names
 307     """
 308     view_layer = bpy.context.view_layer
 309     uarm  = obj.pose.bones[fk[0]]
 310     farm  = obj.pose.bones[fk[1]]
 311     hand  = obj.pose.bones[fk[2]]
 312     uarmi = obj.pose.bones[ik[0]]
 313     farmi = obj.pose.bones[ik[1]]
 314     handi = obj.pose.bones[ik[2]]
 315
 316     if 'auto_stretch' in handi.keys():
 317         # This is kept for compatibility with legacy rigify Human
 318         # Stretch
 319         if handi['auto_stretch'] == 0.0:
 320             uarm['stretch_length'] = handi['stretch_length']
 321         else:
 322             diff = (uarmi.vector.length + farmi.vector.length) / (uarm.vector.length + farm.vector.length)
 323             uarm['stretch_length'] *= diff
 324
 325         # Upper arm position
 326         match_pose_rotation(uarm, uarmi)
 327         match_pose_scale(uarm, uarmi)
 328         view_layer.update()
 329
 330         # Forearm position
 331         match_pose_rotation(farm, farmi)
 332         match_pose_scale(farm, farmi)
 333         view_layer.update()
 334
 335         # Hand position
 336         match_pose_rotation(hand, handi)
 337         match_pose_scale(hand, handi)
 338         view_layer.update()
 339     else:
 340         # Upper arm position
 341         match_pose_translation(uarm, uarmi)
 342         match_pose_rotation(uarm, uarmi)
 343         match_pose_scale(uarm, uarmi)
 344         view_layer.update()
 345
 346         # Forearm position
 347         #match_pose_translation(hand, handi)
 348         match_pose_rotation(farm, farmi)
 349         match_pose_scale(farm, farmi)
 350         view_layer.update()
 351
 352         # Hand position
 353         match_pose_translation(hand, handi)
 354         match_pose_rotation(hand, handi)
 355         match_pose_scale(hand, handi)
 356         view_layer.update()
 357
 358
 359 def ik2fk_arm(obj, fk, ik):
 360     """ Matches the ik bones in an arm rig to the fk bones.
 361         obj: armature object
 362         fk:  list of fk bone names
 363         ik:  list of ik bone names
 364     """
 365     view_layer = bpy.context.view_layer
 366     uarm  = obj.pose.bones[fk[0]]
 367     farm  = obj.pose.bones[fk[1]]
 368     hand  = obj.pose.bones[fk[2]]
 369     uarmi = obj.pose.bones[ik[0]]
 370     farmi = obj.pose.bones[ik[1]]
 371     handi = obj.pose.bones[ik[2]]
 372
 373     main_parent = obj.pose.bones[ik[4]]
 374
 375     if ik[3] != "" and main_parent['pole_vector']:
 376         pole  = obj.pose.bones[ik[3]]
 377     else:
 378         pole = None
 379
 380
 381     if pole:
 382         # Stretch
 383         # handi['stretch_length'] = uarm['stretch_length']
 384
 385         # Hand position
 386         match_pose_translation(handi, hand)
 387         match_pose_rotation(handi, hand)
 388         match_pose_scale(handi, hand)
 389         view_layer.update()
 390
 391         # Pole target position
 392         match_pole_target(view_layer, uarmi, farmi, pole, uarm.matrix, (uarmi.length + farmi.length))
 393
 394     else:
 395         # Hand position
 396         match_pose_translation(handi, hand)
 397         match_pose_rotation(handi, hand)
 398         match_pose_scale(handi, hand)
 399         view_layer.update()
 400
 401         # Upper Arm position
 402         match_pose_translation(uarmi, uarm)
 403         #match_pose_rotation(uarmi, uarm)
 404         set_pose_rotation(uarmi, Matrix())
 405         match_pose_scale(uarmi, uarm)
 406         view_layer.update()
 407
 408         # Rotation Correction
 409         correct_rotation(view_layer, uarmi, uarm.matrix)
 410
 411     correct_scale(view_layer, uarmi, uarm.matrix)
 412 ''']
 413
 414 UTILITIES_FUNC_OLD_LEG_FKIK = ['''
 415 ######################
 416 ## IK Leg functions ##
 417 ######################
 418
 419 def fk2ik_leg(obj, fk, ik):
 420     """ Matches the fk bones in a leg rig to the ik bones.
 421         obj: armature object
 422         fk:  list of fk bone names
 423         ik:  list of ik bone names
 424     """
 425     view_layer = bpy.context.view_layer
 426     thigh  = obj.pose.bones[fk[0]]
 427     shin   = obj.pose.bones[fk[1]]
 428     foot   = obj.pose.bones[fk[2]]
 429     mfoot  = obj.pose.bones[fk[3]]
 430     thighi = obj.pose.bones[ik[0]]
 431     shini  = obj.pose.bones[ik[1]]
 432     footi  = obj.pose.bones[ik[2]]
 433     mfooti = obj.pose.bones[ik[3]]
 434
 435     if 'auto_stretch' in footi.keys():
 436         # This is kept for compatibility with legacy rigify Human
 437         # Stretch
 438         if footi['auto_stretch'] == 0.0:
 439             thigh['stretch_length'] = footi['stretch_length']
 440         else:
 441             diff = (thighi.vector.length + shini.vector.length) / (thigh.vector.length + shin.vector.length)
 442             thigh['stretch_length'] *= diff
 443
 444         # Thigh position
 445         match_pose_rotation(thigh, thighi)
 446         match_pose_scale(thigh, thighi)
 447         view_layer.update()
 448
 449         # Shin position
 450         match_pose_rotation(shin, shini)
 451         match_pose_scale(shin, shini)
 452         view_layer.update()
 453
 454         # Foot position
 455         footmat = get_pose_matrix_in_other_space(mfooti.matrix, foot)
 456         footmat = convert_pose_matrix_via_rest_delta(footmat, mfoot, foot)
 457         set_pose_rotation(foot, footmat)
 458         set_pose_scale(foot, footmat)
 459         view_layer.update()
 460
 461     else:
 462         # Thigh position
 463         match_pose_translation(thigh, thighi)
 464         match_pose_rotation(thigh, thighi)
 465         match_pose_scale(thigh, thighi)
 466         view_layer.update()
 467
 468         # Shin position
 469         match_pose_rotation(shin, shini)
 470         match_pose_scale(shin, shini)
 471         view_layer.update()
 472
 473         # Foot position
 474         footmat = get_pose_matrix_in_other_space(mfooti.matrix, foot)
 475         footmat = convert_pose_matrix_via_rest_delta(footmat, mfoot, foot)
 476         set_pose_rotation(foot, footmat)
 477         set_pose_scale(foot, footmat)
 478         view_layer.update()
 479
 480
 481 def ik2fk_leg(obj, fk, ik):
 482     """ Matches the ik bones in a leg rig to the fk bones.
 483         obj: armature object
 484         fk:  list of fk bone names
 485         ik:  list of ik bone names
 486     """
 487     view_layer = bpy.context.view_layer
 488     thigh    = obj.pose.bones[fk[0]]
 489     shin     = obj.pose.bones[fk[1]]
 490     mfoot    = obj.pose.bones[fk[2]]
 491     if fk[3] != "":
 492         foot      = obj.pose.bones[fk[3]]
 493     else:
 494         foot = None
 495     thighi   = obj.pose.bones[ik[0]]
 496     shini    = obj.pose.bones[ik[1]]
 497     footi    = obj.pose.bones[ik[2]]
 498     footroll = obj.pose.bones[ik[3]]
 499
 500     main_parent = obj.pose.bones[ik[6]]
 501
 502     if ik[4] != "" and main_parent['pole_vector']:
 503         pole     = obj.pose.bones[ik[4]]
 504     else:
 505         pole = None
 506     mfooti   = obj.pose.bones[ik[5]]
 507
 508     if (not pole) and (foot):
 509
 510         # Clear footroll
 511         set_pose_rotation(footroll, Matrix())
 512         view_layer.update()
 513
 514         # Foot position
 515         footmat = get_pose_matrix_in_other_space(foot.matrix, footi)
 516         footmat = convert_pose_matrix_via_rest_delta(footmat, mfooti, footi)
 517         set_pose_translation(footi, footmat)
 518         set_pose_rotation(footi, footmat)
 519         set_pose_scale(footi, footmat)
 520         view_layer.update()
 521
 522         # Thigh position
 523         match_pose_translation(thighi, thigh)
 524         #match_pose_rotation(thighi, thigh)
 525         set_pose_rotation(thighi, Matrix())
 526         match_pose_scale(thighi, thigh)
 527         view_layer.update()
 528
 529         # Rotation Correction
 530         correct_rotation(view_layer, thighi, thigh.matrix)
 531
 532     else:
 533         # Stretch
 534         if 'stretch_length' in footi.keys() and 'stretch_length' in thigh.keys():
 535             # Kept for compat with legacy rigify Human
 536             footi['stretch_length'] = thigh['stretch_length']
 537
 538         # Clear footroll
 539         set_pose_rotation(footroll, Matrix())
 540         view_layer.update()
 541
 542         # Foot position
 543         footmat = get_pose_matrix_in_other_space(mfoot.matrix, footi)
 544         footmat = convert_pose_matrix_via_rest_delta(footmat, mfooti, footi)
 545         set_pose_translation(footi, footmat)
 546         set_pose_rotation(footi, footmat)
 547         set_pose_scale(footi, footmat)
 548         view_layer.update()
 549
 550         # Pole target position
 551         match_pole_target(view_layer, thighi, shini, pole, thigh.matrix, (thighi.length + shini.length))
 552
 553     correct_scale(view_layer, thighi, thigh.matrix)
 554 ''']
 555
 556 UTILITIES_FUNC_OLD_POLE = ['''
 557 ################################
 558 ## IK Rotation-Pole functions ##
 559 ################################
 560
 561 def rotPoleToggle(rig, limb_type, controls, ik_ctrl, fk_ctrl, parent, pole):
 562
 563     rig_id = rig.data['rig_id']
 564     leg_fk2ik = eval('bpy.ops.pose.rigify_leg_fk2ik_' + rig_id)
 565     arm_fk2ik = eval('bpy.ops.pose.rigify_arm_fk2ik_' + rig_id)
 566     leg_ik2fk = eval('bpy.ops.pose.rigify_leg_ik2fk_' + rig_id)
 567     arm_ik2fk = eval('bpy.ops.pose.rigify_arm_ik2fk_' + rig_id)
 568
 569     controls = parse_bone_names(controls)
 570     ik_ctrl = parse_bone_names(ik_ctrl)
 571     fk_ctrl = parse_bone_names(fk_ctrl)
 572     parent = parse_bone_names(parent)
 573     pole = parse_bone_names(pole)
 574
 575     pbones = bpy.context.selected_pose_bones
 576     bpy.ops.pose.select_all(action='DESELECT')
 577
 578     for b in pbones:
 579
 580         new_pole_vector_value = not rig.pose.bones[parent]['pole_vector']
 581
 582         if b.name in controls or b.name in ik_ctrl:
 583             if limb_type == 'arm':
 584                 func1 = arm_fk2ik
 585                 func2 = arm_ik2fk
 586                 rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
 587                 rig.pose.bones[controls[4]].bone.select = not new_pole_vector_value
 588                 rig.pose.bones[parent].bone.select = not new_pole_vector_value
 589                 rig.pose.bones[pole].bone.select = new_pole_vector_value
 590
 591                 kwargs1 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
 592                           'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1],
 593                           'hand_ik': controls[4]}
 594                 kwargs2 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
 595                           'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1], 'hand_ik': controls[4],
 596                           'pole': pole, 'main_parent': parent}
 597             else:
 598                 func1 = leg_fk2ik
 599                 func2 = leg_ik2fk
 600                 rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
 601                 rig.pose.bones[controls[6]].bone.select = not new_pole_vector_value
 602                 rig.pose.bones[controls[5]].bone.select = not new_pole_vector_value
 603                 rig.pose.bones[parent].bone.select = not new_pole_vector_value
 604                 rig.pose.bones[pole].bone.select = new_pole_vector_value
 605
 606                 kwargs1 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
 607                           'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
 608                           'foot_ik': ik_ctrl[2], 'mfoot_ik': ik_ctrl[2]}
 609                 kwargs2 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
 610                           'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
 611                           'foot_ik': controls[6], 'pole': pole, 'footroll': controls[5], 'mfoot_ik': ik_ctrl[2],
 612                           'main_parent': parent}
 613
 614             func1(**kwargs1)
 615             rig.pose.bones[parent]['pole_vector'] = new_pole_vector_value
 616             func2(**kwargs2)
 617
 618             bpy.ops.pose.select_all(action='DESELECT')
 619 ''']
 620
 621 REGISTER_OP_OLD_ARM_FKIK = ['Rigify_Arm_FK2IK', 'Rigify_Arm_IK2FK']
 622
 623 UTILITIES_OP_OLD_ARM_FKIK = ['''
 624 ##################################
 625 ## IK/FK Arm snapping operators ##
 626 ##################################
 627
 628 class Rigify_Arm_FK2IK(bpy.types.Operator):
 629     """ Snaps an FK arm to an IK arm.
 630     """
 631     bl_idname = "pose.rigify_arm_fk2ik_" + rig_id
 632     bl_label = "Rigify Snap FK arm to IK"
 633     bl_options = {'UNDO', 'INTERNAL'}
 634
 635     uarm_fk: StringProperty(name="Upper Arm FK Name")
 636     farm_fk: StringProperty(name="Forerm FK Name")
 637     hand_fk: StringProperty(name="Hand FK Name")
 638
 639     uarm_ik: StringProperty(name="Upper Arm IK Name")
 640     farm_ik: StringProperty(name="Forearm IK Name")
 641     hand_ik: StringProperty(name="Hand IK Name")
 642
 643     @classmethod
 644     def poll(cls, context):
 645         return (context.active_object != None and context.mode == 'POSE')
 646
 647     def execute(self, context):
 648         fk2ik_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
 649         return {'FINISHED'}
 650
 651
 652 class Rigify_Arm_IK2FK(bpy.types.Operator):
 653     """ Snaps an IK arm to an FK arm.
 654     """
 655     bl_idname = "pose.rigify_arm_ik2fk_" + rig_id
 656     bl_label = "Rigify Snap IK arm to FK"
 657     bl_options = {'UNDO', 'INTERNAL'}
 658
 659     uarm_fk: StringProperty(name="Upper Arm FK Name")
 660     farm_fk: StringProperty(name="Forerm FK Name")
 661     hand_fk: StringProperty(name="Hand FK Name")
 662
 663     uarm_ik: StringProperty(name="Upper Arm IK Name")
 664     farm_ik: StringProperty(name="Forearm IK Name")
 665     hand_ik: StringProperty(name="Hand IK Name")
 666     pole   : StringProperty(name="Pole IK Name")
 667
 668     main_parent: StringProperty(name="Main Parent", default="")
 669
 670     @classmethod
 671     def poll(cls, context):
 672         return (context.active_object != None and context.mode == 'POSE')
 673
 674     def execute(self, context):
 675         ik2fk_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole, self.main_parent])
 676         return {'FINISHED'}
 677 ''']
 678
 679 REGISTER_OP_OLD_LEG_FKIK = ['Rigify_Leg_FK2IK', 'Rigify_Leg_IK2FK']
 680
 681 UTILITIES_OP_OLD_LEG_FKIK = ['''
 682 ##################################
 683 ## IK/FK Leg snapping operators ##
 684 ##################################
 685
 686 class Rigify_Leg_FK2IK(bpy.types.Operator):
 687     """ Snaps an FK leg to an IK leg.
 688     """
 689     bl_idname = "pose.rigify_leg_fk2ik_" + rig_id
 690     bl_label = "Rigify Snap FK leg to IK"
 691     bl_options = {'UNDO', 'INTERNAL'}
 692
 693     thigh_fk: StringProperty(name="Thigh FK Name")
 694     shin_fk:  StringProperty(name="Shin FK Name")
 695     foot_fk:  StringProperty(name="Foot FK Name")
 696     mfoot_fk: StringProperty(name="MFoot FK Name")
 697
 698     thigh_ik: StringProperty(name="Thigh IK Name")
 699     shin_ik:  StringProperty(name="Shin IK Name")
 700     foot_ik:  StringProperty(name="Foot IK Name")
 701     mfoot_ik: StringProperty(name="MFoot IK Name")
 702
 703     @classmethod
 704     def poll(cls, context):
 705         return (context.active_object != None and context.mode == 'POSE')
 706
 707     def execute(self, context):
 708         fk2ik_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.foot_fk, self.mfoot_fk], ik=[self.thigh_ik, self.shin_ik, self.foot_ik, self.mfoot_ik])
 709         return {'FINISHED'}
 710
 711
 712 class Rigify_Leg_IK2FK(bpy.types.Operator):
 713     """ Snaps an IK leg to an FK leg.
 714     """
 715     bl_idname = "pose.rigify_leg_ik2fk_" + rig_id
 716     bl_label = "Rigify Snap IK leg to FK"
 717     bl_options = {'UNDO', 'INTERNAL'}
 718
 719     thigh_fk: StringProperty(name="Thigh FK Name")
 720     shin_fk:  StringProperty(name="Shin FK Name")
 721     mfoot_fk: StringProperty(name="MFoot FK Name")
 722     foot_fk:  StringProperty(name="Foot FK Name", default="")
 723     thigh_ik: StringProperty(name="Thigh IK Name")
 724     shin_ik:  StringProperty(name="Shin IK Name")
 725     foot_ik:  StringProperty(name="Foot IK Name")
 726     footroll: StringProperty(name="Foot Roll Name")
 727     pole:     StringProperty(name="Pole IK Name")
 728     mfoot_ik: StringProperty(name="MFoot IK Name")
 729
 730     main_parent: StringProperty(name="Main Parent", default="")
 731
 732     @classmethod
 733     def poll(cls, context):
 734         return (context.active_object != None and context.mode == 'POSE')
 735
 736     def execute(self, context):
 737         ik2fk_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.mfoot_fk, self.foot_fk], ik=[self.thigh_ik, self.shin_ik, self.foot_ik, self.footroll, self.pole, self.mfoot_ik, self.main_parent])
 738         return {'FINISHED'}
 739 ''']
 740
 741 REGISTER_OP_OLD_POLE = ['Rigify_Rot2PoleSwitch']
 742
 743 UTILITIES_OP_OLD_POLE = ['''
 744 ###########################
 745 ## IK Rotation Pole Snap ##
 746 ###########################
 747
 748 class Rigify_Rot2PoleSwitch(bpy.types.Operator):
 749     bl_idname = "pose.rigify_rot2pole_" + rig_id
 750     bl_label = "Rotation - Pole toggle"
 751     bl_description = "Toggles IK chain between rotation and pole target"
 752
 753     bone_name: StringProperty(default='')
 754     limb_type: StringProperty(name="Limb Type")
 755     controls: StringProperty(name="Controls string")
 756     ik_ctrl: StringProperty(name="IK Controls string")
 757     fk_ctrl: StringProperty(name="FK Controls string")
 758     parent: StringProperty(name="Parent name")
 759     pole: StringProperty(name="Pole name")
 760
 761     def execute(self, context):
 762         rig = context.object
 763
 764         if self.bone_name:
 765             bpy.ops.pose.select_all(action='DESELECT')
 766             rig.pose.bones[self.bone_name].bone.select = True
 767
 768         rotPoleToggle(rig, self.limb_type, self.controls, self.ik_ctrl, self.fk_ctrl, self.parent, self.pole)
 769         return {'FINISHED'}
 770 ''']
 771
 772 REGISTER_RIG_OLD_ARM = REGISTER_OP_OLD_ARM_FKIK + REGISTER_OP_OLD_POLE
 773
 774 UTILITIES_RIG_OLD_ARM = [
 775     *UTILITIES_FUNC_COMMON_IKFK,
 776     *UTILITIES_FUNC_OLD_ARM_FKIK,
 777     *UTILITIES_FUNC_OLD_POLE,
 778     *UTILITIES_OP_OLD_ARM_FKIK,
 779     *UTILITIES_OP_OLD_POLE,
 780 ]
 781
 782 REGISTER_RIG_OLD_LEG = REGISTER_OP_OLD_LEG_FKIK + REGISTER_OP_OLD_POLE
 783
 784 UTILITIES_RIG_OLD_LEG = [
 785     *UTILITIES_FUNC_COMMON_IKFK,
 786     *UTILITIES_FUNC_OLD_LEG_FKIK,
 787     *UTILITIES_FUNC_OLD_POLE,
 788     *UTILITIES_OP_OLD_LEG_FKIK,
 789     *UTILITIES_OP_OLD_POLE,
 790 ]
 791
 792 ##############################
 793 ## Default set of utilities ##
 794 ##############################
 795
 796 UI_REGISTER = [
 797     'RigUI',
 798     'RigLayers',
 799 ]
 800
 801 UI_UTILITIES = [
 802 ]
 803
 804 UI_SLIDERS = '''
 805 ###################
 806 ## Rig UI Panels ##
 807 ###################
 808
 809 class RigUI(bpy.types.Panel):
 810     bl_space_type = 'VIEW_3D'
 811     bl_region_type = 'UI'
 812     bl_label = "Rig Main Properties"
 813     bl_idname = "VIEW3D_PT_rig_ui_" + rig_id
 814     bl_category = 'Item'
 815
 816     @classmethod
 817     def poll(self, context):
 818         if context.mode != 'POSE':
 819             return False
 820         try:
 821             return (context.active_object.data.get("rig_id") == rig_id)
 822         except (AttributeError, KeyError, TypeError):
 823             return False
 824
 825     def draw(self, context):
 826         layout = self.layout
 827         pose_bones = context.active_object.pose.bones
 828         try:
 829             selected_bones = set(bone.name for bone in context.selected_pose_bones)
 830             selected_bones.add(context.active_pose_bone.name)
 831         except (AttributeError, TypeError):
 832             return
 833
 834         def is_selected(names):
 835             # Returns whether any of the named bones are selected.
 836             if isinstance(names, list) or isinstance(names, set):
 837                 return not selected_bones.isdisjoint(names)
 838             elif names in selected_bones:
 839                 return True
 840             return False
 841
 842         num_rig_separators = [-1]
 843
 844         def emit_rig_separator():
 845             if num_rig_separators[0] >= 0:
 846                 layout.separator()
 847             num_rig_separators[0] += 1
 848 '''
 849
 850 UI_REGISTER_BAKE_SETTINGS = ['RigBakeSettings']
 851
 852 UI_BAKE_SETTINGS = '''
 853 class RigBakeSettings(bpy.types.Panel):
 854     bl_space_type = 'VIEW_3D'
 855     bl_region_type = 'UI'
 856     bl_label = "Rig Bake Settings"
 857     bl_idname = "VIEW3D_PT_rig_bake_settings_" + rig_id
 858     bl_category = 'Item'
 859
 860     @classmethod
 861     def poll(self, context):
 862         return RigUI.poll(context) and find_action(context.active_object) is not None
 863
 864     def draw(self, context):
 865         RigifyBakeKeyframesMixin.draw_common_bake_ui(context, self.layout)
 866 '''
 867
 868 def layers_ui(layers, layout):
 869     """ Turn a list of booleans + a list of names into a layer UI.
 870     """
 871
 872     code = '''
 873 class RigLayers(bpy.types.Panel):
 874     bl_space_type = 'VIEW_3D'
 875     bl_region_type = 'UI'
 876     bl_label = "Rig Layers"
 877     bl_idname = "VIEW3D_PT_rig_layers_" + rig_id
 878     bl_category = 'Item'
 879
 880     @classmethod
 881     def poll(self, context):
 882         try:
 883             return (context.active_object.data.get("rig_id") == rig_id)
 884         except (AttributeError, KeyError, TypeError):
 885             return False
 886
 887     def draw(self, context):
 888         layout = self.layout
 889         col = layout.column()
 890 '''
 891     rows = {}
 892     for i in range(28):
 893         if layers[i]:
 894             if layout[i][1] not in rows:
 895                 rows[layout[i][1]] = []
 896             rows[layout[i][1]] += [(layout[i][0], i)]
 897
 898     keys = list(rows.keys())
 899     keys.sort()
 900
 901     for key in keys:
 902         code += "\n        row = col.row()\n"
 903         i = 0
 904         for l in rows[key]:
 905             if i > 3:
 906                 code += "\n        row = col.row()\n"
 907                 i = 0
 908             code += "        row.prop(context.active_object.data, 'layers', index=%s, toggle=True, text='%s')\n" % (str(l[1]), l[0])
 909             i += 1
 910
 911     # Root layer
 912     code += "\n        row = col.row()"
 913     code += "\n        row.separator()"
 914     code += "\n        row = col.row()"
 915     code += "\n        row.separator()\n"
 916     code += "\n        row = col.row()\n"
 917     code += "        row.prop(context.active_object.data, 'layers', index=28, toggle=True, text='Root')\n"
 918
 919     return code
 920
 921
 922 def quote_parameters(positional, named):
 923     """Quote the given positional and named parameters as a code string."""
 924     positional_list = [ repr(v) for v in positional ]
 925     named_list = [ "%s=%r" % (k, v) for k, v in named.items() ]
 926     return ', '.join(positional_list + named_list)
 927
 928 def indent_lines(lines, indent=4):
 929     if indent > 0:
 930         prefix = ' ' * indent
 931         return [ prefix + line for line in lines ]
 932     else:
 933         return lines
 934
 935
 936 class PanelLayout(object):
 937     """Utility class that builds code for creating a layout."""
 938
 939     def __init__(self, parent, index=0):
 940         if isinstance(parent, PanelLayout):
 941             self.parent = parent
 942             self.script = parent.script
 943         else:
 944             self.parent = None
 945             self.script = parent
 946
 947         self.header = []
 948         self.items = []
 949         self.indent = 0
 950         self.index = index
 951         self.layout = self._get_layout_var(index)
 952         self.is_empty = True
 953
 954     @staticmethod
 955     def _get_layout_var(index):
 956         return 'layout' if index == 0 else 'group' + str(index)
 957
 958     def clear_empty(self):
 959         self.is_empty = False
 960
 961         if self.parent:
 962             self.parent.clear_empty()
 963
 964     def get_lines(self):
 965         lines = []
 966
 967         for item in self.items:
 968             if isinstance(item, PanelLayout):
 969                 lines += item.get_lines()
 970             else:
 971                 lines.append(item)
 972
 973         if len(lines) > 0:
 974             return self.wrap_lines(lines)
 975         else:
 976             return []
 977
 978     def wrap_lines(self, lines):
 979         return self.header + indent_lines(lines, self.indent)
 980
 981     def add_line(self, line):
 982         assert isinstance(line, str)
 983
 984         self.items.append(line)
 985
 986         if self.is_empty:
 987             self.clear_empty()
 988
 989     def use_bake_settings(self):
 990         """This panel contains operators that need the common Bake settings."""
 991         self.parent.use_bake_settings()
 992
 993     def custom_prop(self, bone_name, prop_name, **params):
 994         """Add a custom property input field to the panel."""
 995         param_str = quote_parameters([ rna_idprop_quote_path(prop_name) ], params)
 996         self.add_line(
 997             "%s.prop(pose_bones[%r], %s)" % (self.layout, bone_name, param_str)
 998         )
 999
1000     def operator(self, operator_name, *, properties=None, **params):
1001         """Add an operator call button to the panel."""
1002         name = operator_name.format_map(self.script.format_args)
1003         param_str = quote_parameters([ name ], params)
1004         call_str = "%s.operator(%s)" % (self.layout, param_str)
1005         if properties:
1006             self.add_line("props = " + call_str)
1007             for k, v in properties.items():
1008                 self.add_line("props.%s = %r" % (k,v))
1009         else:
1010             self.add_line(call_str)
1011
1012     def add_nested_layout(self, name, params):
1013         param_str = quote_parameters([], params)
1014         sub_panel = PanelLayout(self, self.index + 1)
1015         sub_panel.header.append('%s = %s.%s(%s)' % (sub_panel.layout, self.layout, name, param_str))
1016         self.items.append(sub_panel)
1017         return sub_panel
1018
1019     def row(self, **params):
1020         """Add a nested row layout to the panel."""
1021         return self.add_nested_layout('row', params)
1022
1023     def column(self, **params):
1024         """Add a nested column layout to the panel."""
1025         return self.add_nested_layout('column', params)
1026
1027     def split(self, **params):
1028         """Add a split layout to the panel."""
1029         return self.add_nested_layout('split', params)
1030
1031
1032 class BoneSetPanelLayout(PanelLayout):
1033     """Panel restricted to a certain set of bones."""
1034
1035     def __init__(self, rig_panel, bones):
1036         assert isinstance(bones, frozenset)
1037         super().__init__(rig_panel)
1038         self.bones = bones
1039         self.show_bake_settings = False
1040
1041     def clear_empty(self):
1042         self.parent.bones |= self.bones
1043
1044         super().clear_empty()
1045
1046     def wrap_lines(self, lines):
1047         if self.bones != self.parent.bones:
1048             header = ["if is_selected(%r):" % (set(self.bones))]
1049             return header + indent_lines(lines)
1050         else:
1051             return lines
1052
1053     def use_bake_settings(self):
1054         self.show_bake_settings = True
1055         if not self.script.use_bake_settings:
1056             self.script.use_bake_settings = True
1057             self.script.add_utilities(SCRIPT_UTILITIES_BAKE)
1058             self.script.register_classes(SCRIPT_REGISTER_BAKE)
1059
1060
1061 class RigPanelLayout(PanelLayout):
1062     """Panel owned by a certain rig."""
1063
1064     def __init__(self, script, rig):
1065         super().__init__(script)
1066         self.bones = set()
1067         self.subpanels = OrderedDict()
1068
1069     def wrap_lines(self, lines):
1070         header = [ "if is_selected(%r):" % (set(self.bones)) ]
1071         prefix = [ "emit_rig_separator()" ]
1072         return header + indent_lines(prefix + lines)
1073
1074     def panel_with_selected_check(self, control_names):
1075         selected_set = frozenset(control_names)
1076
1077         if selected_set in self.subpanels:
1078             return self.subpanels[selected_set]
1079         else:
1080             panel = BoneSetPanelLayout(self, selected_set)
1081             self.subpanels[selected_set] = panel
1082             self.items.append(panel)
1083             return panel
1084
1085
1086 class ScriptGenerator(base_generate.GeneratorPlugin):
1087     """Generator plugin that builds the python script attached to the rig."""
1088
1089     priority = -100
1090
1091     def __init__(self, generator):
1092         super().__init__(generator)
1093
1094         self.ui_scripts = []
1095         self.ui_imports = UI_IMPORTS.copy()
1096         self.ui_utilities = UI_UTILITIES.copy()
1097         self.ui_register = UI_REGISTER.copy()
1098         self.ui_register_drivers = []
1099         self.ui_register_props = []
1100
1101         self.ui_rig_panels = OrderedDict()
1102
1103         self.use_bake_settings = False
1104
1105     # Structured panel code generation
1106     def panel_with_selected_check(self, rig, control_names):
1107         """Add a panel section with restricted selection."""
1108         rig_key = id(rig)
1109
1110         if rig_key in self.ui_rig_panels:
1111             panel = self.ui_rig_panels[rig_key]
1112         else:
1113             panel = RigPanelLayout(self, rig)
1114             self.ui_rig_panels[rig_key] = panel
1115
1116         return panel.panel_with_selected_check(control_names)
1117
1118     # Raw output
1119     def add_panel_code(self, str_list):
1120         """Add raw code to the panel."""
1121         self.ui_scripts += str_list
1122
1123     def add_imports(self, str_list):
1124         self.ui_imports += str_list
1125
1126     def add_utilities(self, str_list):
1127         self.ui_utilities += str_list
1128
1129     def register_classes(self, str_list):
1130         self.ui_register += str_list
1131
1132     def register_driver_functions(self, str_list):
1133         self.ui_register_drivers += str_list
1134
1135     def register_property(self, name, definition):
1136         self.ui_register_props.append((name, definition))
1137
1138     def initialize(self):
1139         self.format_args = {
1140             'rig_id': self.generator.rig_id,
1141         }
1142
1143     def finalize(self):
1144         metarig = self.generator.metarig
1145         rig_id = self.generator.rig_id
1146
1147         vis_layers = self.obj.data.layers
1148
1149         # Ensure the collection of layer names exists
1150         for i in range(1 + len(metarig.data.rigify_layers), 29):
1151             metarig.data.rigify_layers.add()
1152
1153         # Create list of layer name/row pairs
1154         layer_layout = []
1155         for l in metarig.data.rigify_layers:
1156             layer_layout += [(l.name, l.row)]
1157
1158         # Generate the UI script
1159         script = metarig.data.rigify_rig_ui
1160
1161         if script:
1162             script.clear()
1163         else:
1164             script = bpy.data.texts.new("rig_ui.py")
1165             metarig.data.rigify_rig_ui = script
1166
1167         for s in OrderedDict.fromkeys(self.ui_imports):
1168             script.write(s + "\n")
1169
1170         script.write(UI_BASE_UTILITIES % rig_id)
1171
1172         for s in OrderedDict.fromkeys(self.ui_utilities):
1173             script.write(s + "\n")
1174
1175         script.write(UI_SLIDERS)
1176
1177         for s in self.ui_scripts:
1178             script.write("\n        " + s.replace("\n", "\n        ") + "\n")
1179
1180         if len(self.ui_scripts) > 0:
1181             script.write("\n        num_rig_separators[0] = 0\n")
1182
1183         for panel in self.ui_rig_panels.values():
1184             lines = panel.get_lines()
1185             if len(lines) > 1:
1186                 script.write("\n        ".join([''] + lines) + "\n")
1187
1188         if self.use_bake_settings:
1189             self.ui_register = UI_REGISTER_BAKE_SETTINGS + self.ui_register
1190             script.write(UI_BAKE_SETTINGS)
1191
1192         script.write(layers_ui(vis_layers, layer_layout))
1193
1194         script.write("\ndef register():\n")
1195
1196         ui_register = OrderedDict.fromkeys(self.ui_register)
1197         for s in ui_register:
1198             script.write("    bpy.utils.register_class("+s+")\n")
1199
1200         ui_register_drivers = OrderedDict.fromkeys(self.ui_register_drivers)
1201         for s in ui_register_drivers:
1202             script.write("    bpy.app.driver_namespace['"+s+"'] = "+s+"\n")
1203
1204         ui_register_props = OrderedDict.fromkeys(self.ui_register_props)
1205         for s in ui_register_props:
1206             script.write("    bpy.types.%s = %s\n " % (*s,))
1207
1208         script.write("\ndef unregister():\n")
1209
1210         for s in ui_register_props:
1211             script.write("    del bpy.types.%s\n" % s[0])
1212
1213         for s in ui_register:
1214             script.write("    bpy.utils.unregister_class("+s+")\n")
1215
1216         for s in ui_register_drivers:
1217             script.write("    del bpy.app.driver_namespace['"+s+"']\n")
1218
1219         script.write("\nregister()\n")
1220         script.use_module = True
1221
1222         # Run UI script
1223         exec(script.as_string(), {})
1224
1225         # Attach the script to the rig
1226         self.obj['rig_ui'] = script