rigify/rig_ui_template.py

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