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