rigify/rig_ui_template.py

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