rigify/utils/mechanism.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 import re
   7
   8 from typing import TYPE_CHECKING, Optional, Any, Sequence, Iterable
   9
  10 from bpy.types import (bpy_prop_collection, Material, Object, PoseBone, Driver, FCurve,
  11                        DriverTarget, ID, bpy_struct, FModifierGenerator, Constraint, AnimData,
  12                        ArmatureConstraint)
  13
  14 from rna_prop_ui import rna_idprop_ui_create
  15 from rna_prop_ui import rna_idprop_quote_path as quote_property
  16
  17 from .misc import force_lazy, ArmatureObject, Lazy, OptionalLazy
  18
  19 if TYPE_CHECKING:
  20     from ..base_rig import BaseRig
  21
  22
  23 ##############################################
  24 # Constraint creation utilities
  25 ##############################################
  26
  27 _TRACK_AXIS_MAP = {
  28     'X': 'TRACK_X', '-X': 'TRACK_NEGATIVE_X',
  29     'Y': 'TRACK_Y', '-Y': 'TRACK_NEGATIVE_Y',
  30     'Z': 'TRACK_Z', '-Z': 'TRACK_NEGATIVE_Z',
  31 }
  32
  33
  34 def _set_default_attr(obj, options, attr, value):
  35     if hasattr(obj, attr):
  36         options.setdefault(attr, value)
  37
  38
  39 def make_constraint(
  40         owner: Object | PoseBone, con_type: str,
  41         target: Optional[Object] = None,
  42         subtarget: OptionalLazy[str] = None, *,
  43         insert_index: Optional[int] = None,
  44         space: Optional[str] = None,
  45         track_axis: Optional[str] = None,
  46         use_xyz: Optional[Sequence[bool]] = None,
  47         use_limit_xyz: Optional[Sequence[bool]] = None,
  48         invert_xyz: Optional[Sequence[bool]] = None,
  49         targets: Optional[list[Lazy[str | tuple | dict]]] = None,
  50         **options):
  51     """
  52     Creates and initializes constraint of the specified type for the owner bone.
  53
  54     Specially handled keyword arguments:
  55
  56       target, subtarget: if both not None, passed through to the constraint
  57       insert_index     : insert at the specified index in the stack, instead of at the end
  58       space            : assigned to both owner_space and target_space
  59       track_axis       : allows shorter X, Y, Z, -X, -Y, -Z notation
  60       use_xyz          : list of 3 items is assigned to use_x, use_y and use_z options
  61       use_limit_xyz    : list of 3 items is assigned to use_limit_x/y/z options
  62       invert_xyz       : list of 3 items is assigned to invert_x, invert_y and invert_z options
  63       min/max_x/y/z    : a corresponding use_(min/max/limit)_(x/y/z) option is set to True
  64       targets          : list of strings, tuples or dicts describing Armature constraint targets
  65
  66     Other keyword arguments are directly assigned to the constraint options.
  67     Returns the newly created constraint.
  68
  69     Target bone names can be provided via 'lazy' callable closures without arguments.
  70     """
  71     con = owner.constraints.new(con_type)
  72
  73     # For Armature constraints, allow passing a "targets" list as a keyword argument.
  74     if targets is not None:
  75         assert isinstance(con, ArmatureConstraint)
  76         for target_info in targets:
  77             con_target = con.targets.new()
  78             con_target.target = owner.id_data
  79             # List element can be a string, a tuple or a dictionary.
  80             target_info = force_lazy(target_info)
  81             if isinstance(target_info, str):
  82                 con_target.subtarget = target_info
  83             elif isinstance(target_info, tuple):
  84                 if len(target_info) == 2:
  85                     con_target.subtarget, con_target.weight = map(force_lazy, target_info)
  86                 else:
  87                     con_target.target, con_target.subtarget, con_target.weight = map(force_lazy, target_info)
  88             else:
  89                 assert isinstance(target_info, dict)
  90                 for key, val in target_info.items():
  91                     setattr(con_target, key, force_lazy(val))
  92
  93     if insert_index is not None:
  94         owner.constraints.move(len(owner.constraints)-1, insert_index)
  95
  96     if target is not None and hasattr(con, 'target'):
  97         con.target = target
  98
  99     if subtarget is not None:
 100         con.subtarget = force_lazy(subtarget)
 101
 102     if space is not None:
 103         _set_default_attr(con, options, 'owner_space', space)
 104         _set_default_attr(con, options, 'target_space', space)
 105
 106     if track_axis is not None:
 107         con.track_axis = _TRACK_AXIS_MAP.get(track_axis, track_axis)
 108
 109     if use_xyz is not None:
 110         con.use_x, con.use_y, con.use_z = use_xyz[0:3]
 111
 112     if use_limit_xyz is not None:
 113         con.use_limit_x, con.use_limit_y, con.use_limit_z = use_limit_xyz[0:3]
 114
 115     if invert_xyz is not None:
 116         con.invert_x, con.invert_y, con.invert_z = invert_xyz[0:3]
 117
 118     for key in ['min_x', 'max_x', 'min_y', 'max_y', 'min_z', 'max_z']:
 119         if key in options:
 120             _set_default_attr(con, options, 'use_'+key, True)
 121             _set_default_attr(con, options, 'use_limit_'+key[-1], True)
 122
 123     for p, v in options.items():
 124         setattr(con, p, force_lazy(v))
 125
 126     return con
 127
 128
 129 ##############################################
 130 # Custom property creation utilities
 131 ##############################################
 132
 133 # noinspection PyShadowingBuiltins
 134 def make_property(
 135         owner: bpy_struct, name: str, default, *,
 136         min: float = 0, max: float = 1, soft_min=None, soft_max=None,
 137         description: Optional[str] = None, overridable=True, subtype: Optional[str] = None,
 138         **options):
 139     """
 140     Creates and initializes a custom property of owner.
 141
 142     The soft_min and soft_max parameters default to min and max.
 143     Description defaults to the property name.
 144     """
 145
 146     # Some keyword argument defaults differ
 147     rna_idprop_ui_create(
 148         owner, name, default=default,
 149         min=min, max=max, soft_min=soft_min, soft_max=soft_max,
 150         description=description or name,
 151         overridable=overridable,
 152         subtype=subtype,
 153         **options  # noqa
 154     )
 155
 156
 157 ##############################################
 158 # Driver creation utilities
 159 ##############################################
 160
 161 def _init_driver_target(drv_target: DriverTarget, var_info, target_id: Optional[ID]):
 162     """Initialize a driver variable target from a specification."""
 163
 164     # Parse the simple list format for the common case.
 165     if isinstance(var_info, tuple):
 166         # [ (target_id,) subtarget, ...path ]
 167
 168         # If target_id is supplied as parameter, allow omitting it
 169         if target_id is None or isinstance(var_info[0], bpy.types.ID):
 170             target_id, subtarget, *refs = var_info
 171         else:
 172             subtarget, *refs = var_info
 173
 174         subtarget = force_lazy(subtarget)
 175
 176         # Simple path string case.
 177         if len(refs) == 0:
 178             # [ (target_id,) path_str ]
 179             path = subtarget
 180         else:
 181             # If subtarget is a string, look up a bone in the target
 182             if isinstance(subtarget, str):
 183                 subtarget = target_id.pose.bones[subtarget]
 184
 185             if subtarget == target_id:
 186                 path = ''
 187             else:
 188                 path = subtarget.path_from_id()
 189
 190             # Use ".foo" type path items verbatim, otherwise quote
 191             for item in refs:
 192                 if isinstance(item, str):
 193                     path += item if item[0] == '.' else quote_property(item)
 194                 else:
 195                     path += f'[{repr(item)}]'
 196
 197             if path[0] == '.':
 198                 path = path[1:]
 199
 200         drv_target.id = target_id
 201         drv_target.data_path = path
 202
 203     else:
 204         # { 'id': ..., ... }
 205         target_id = var_info.get('id', target_id)
 206
 207         if target_id is not None:
 208             drv_target.id = target_id
 209
 210         for tp, tv in var_info.items():
 211             setattr(drv_target, tp, force_lazy(tv))
 212
 213
 214 def _add_driver_variable(drv: Driver, var_name: str, var_info, target_id: Optional[ID]):
 215     """Add and initialize a driver variable."""
 216
 217     var = drv.variables.new()
 218     var.name = var_name
 219
 220     # Parse the simple list format for the common case.
 221     if isinstance(var_info, tuple):
 222         # [ (target_id,) subtarget, ...path ]
 223         var.type = "SINGLE_PROP"
 224
 225         _init_driver_target(var.targets[0], var_info, target_id)
 226
 227     else:
 228         # Variable info as generic dictionary - assign properties.
 229         # { 'type': 'SINGLE_PROP', 'targets':[...] }
 230         var.type = var_info['type']
 231
 232         for p, v in var_info.items():
 233             if p == 'targets':
 234                 for i, tdata in enumerate(v):
 235                     _init_driver_target(var.targets[i], tdata, target_id)
 236             elif p != 'type':
 237                 setattr(var, p, force_lazy(v))
 238
 239
 240 # noinspection PyShadowingBuiltins
 241 def make_driver(owner: bpy_struct, prop: str, *, index=-1, type='SUM',
 242                 expression: Optional[str] = None,
 243                 variables: Iterable | dict = (),
 244                 polynomial: Optional[list[float]] = None,
 245                 target_id: Optional[ID] = None) -> FCurve:
 246     """
 247     Creates and initializes a driver for the 'prop' property of owner.
 248
 249     Arguments:
 250       owner           : object to add the driver to
 251       prop            : property of the object to add the driver to
 252       index           : item index for vector properties
 253       type            : built-in driver math operation (incompatible with expression)
 254       expression      : custom driver expression
 255       variables       : either a list or dictionary of variable specifications.
 256       polynomial      : coefficients of the POLYNOMIAL driver modifier
 257       target_id       : specifies the target ID of variables implicitly
 258
 259     Specification format:
 260         If the variables argument is a dictionary, keys specify variable names.
 261         Otherwise, names are set to var, var1, var2, ... etc.:
 262
 263           variables = [ ..., ..., ... ]
 264           variables = { 'var': ..., 'var1': ..., 'var2': ... }
 265
 266         Variable specifications are constructed as nested dictionaries and lists that
 267         follow the property structure of the original Blender objects, but the most
 268         common case can be abbreviated as a simple tuple.
 269
 270         The following specifications are equivalent:
 271
 272           ( target, subtarget, '.foo', 'bar' )
 273
 274           { 'type': 'SINGLE_PROP', 'targets':[( target, subtarget, '.foo', 'bar' )] }
 275
 276           { 'type': 'SINGLE_PROP',
 277             'targets':[{ 'id': target, 'data_path': subtarget.path_from_id() + '.foo["bar"]' }] }
 278
 279         If subtarget is as string, it is automatically looked up within target as a bone.
 280
 281         It is possible to specify path directly as a simple string without following items:
 282
 283           ( target, 'path' )
 284
 285           { 'type': 'SINGLE_PROP', 'targets':[{ 'id': target, 'data_path': 'path' }] }
 286
 287         If the target_id parameter is not None, it is possible to omit target:
 288
 289           ( subtarget, '.foo', 'bar' )
 290
 291           { 'type': 'SINGLE_PROP',
 292             'targets':[{ 'id': target_id, 'data_path': subtarget.path_from_id() + '.foo["bar"]' }] }
 293
 294     Returns the newly created driver FCurve.
 295
 296     Target bone names can be provided via 'lazy' callable closures without arguments.
 297     """
 298     fcu = owner.driver_add(prop, index)
 299     drv = fcu.driver
 300
 301     if expression is not None:
 302         drv.type = 'SCRIPTED'
 303         drv.expression = expression
 304     else:
 305         drv.type = type
 306
 307     # In case the driver already existed, remove contents
 308     for var in list(drv.variables):
 309         drv.variables.remove(var)
 310
 311     for mod in list(fcu.modifiers):
 312         fcu.modifiers.remove(mod)
 313
 314     # Fill in new data
 315     if not isinstance(variables, dict):
 316         # variables = [ info, ... ]
 317         for i, var_info in enumerate(variables):
 318             var_name = 'var' if i == 0 else 'var' + str(i)
 319             _add_driver_variable(drv, var_name, var_info, target_id)
 320     else:
 321         # variables = { 'varname': info, ... }
 322         for var_name, var_info in variables.items():
 323             _add_driver_variable(drv, var_name, var_info, target_id)
 324
 325     if polynomial is not None:
 326         drv_modifier = fcu.modifiers.new('GENERATOR')
 327         assert isinstance(drv_modifier, FModifierGenerator)
 328         drv_modifier.mode = 'POLYNOMIAL'
 329         drv_modifier.poly_order = len(polynomial)-1
 330         for i, v in enumerate(polynomial):
 331             drv_modifier.coefficients[i] = v
 332
 333     return fcu
 334
 335
 336 ##############################################
 337 # Driver variable utilities
 338 ##############################################
 339
 340 # noinspection PyShadowingBuiltins
 341 def driver_var_transform(target: ID, bone: Optional[str] = None, *,
 342                          type='LOC_X', space='WORLD', rotation_mode='AUTO'):
 343     """
 344     Create a Transform Channel driver variable specification.
 345
 346     Usage:
 347         make_driver(..., variables=[driver_var_transform(...)])
 348
 349     Target bone name can be provided via a 'lazy' callable closure without arguments.
 350     """
 351
 352     assert space in {'WORLD', 'TRANSFORM', 'LOCAL'}
 353
 354     target_map = {
 355         'id': target,
 356         'transform_type': type,
 357         'transform_space': space + '_SPACE',
 358         'rotation_mode': rotation_mode,
 359     }
 360
 361     if bone is not None:
 362         target_map['bone_target'] = bone
 363
 364     return {'type': 'TRANSFORMS', 'targets': [target_map]}
 365
 366
 367 def driver_var_distance(target: ID, *,
 368                         bone1: Optional[str] = None,
 369                         target2: Optional[ID] = None,
 370                         bone2: Optional[str] = None,
 371                         space1='WORLD', space2='WORLD'):
 372     """
 373     Create a Distance driver variable specification.
 374
 375     Usage:
 376         make_driver(..., variables=[driver_var_distance(...)])
 377
 378     Target bone name can be provided via a 'lazy' callable closure without arguments.
 379     """
 380
 381     assert space1 in {'WORLD', 'TRANSFORM', 'LOCAL'}
 382     assert space2 in {'WORLD', 'TRANSFORM', 'LOCAL'}
 383
 384     target1_map = {
 385         'id': target,
 386         'transform_space': space1 + '_SPACE',
 387     }
 388
 389     if bone1 is not None:
 390         target1_map['bone_target'] = bone1
 391
 392     target2_map = {
 393         'id': target2 or target,
 394         'transform_space': space2 + '_SPACE',
 395     }
 396
 397     if bone2 is not None:
 398         target2_map['bone_target'] = bone2
 399
 400     return {'type': 'LOC_DIFF', 'targets': [target1_map, target2_map]}
 401
 402
 403 ##############################################
 404 # Constraint management
 405 ##############################################
 406
 407 def move_constraint(source: Object | PoseBone, target: Object | PoseBone | str, con: Constraint):
 408     """
 409     Move a constraint from one owner to another, together with drivers.
 410     """
 411
 412     assert source.constraints[con.name] == con
 413
 414     if isinstance(target, str):
 415         target = con.id_data.pose.bones[target]
 416
 417     con_tgt = target.constraints.copy(con)
 418
 419     if target.id_data == con.id_data:
 420         adt = con.id_data.animation_data
 421         if adt:
 422             prefix = con.path_from_id()
 423             new_prefix = con_tgt.path_from_id()
 424             for fcu in adt.drivers:
 425                 if fcu.data_path.startswith(prefix):
 426                     fcu.data_path = new_prefix + fcu.data_path[len(prefix):]
 427
 428     source.constraints.remove(con)
 429
 430
 431 def move_all_constraints(obj: Object,
 432                          source: Object | PoseBone | str,
 433                          target: Object | PoseBone | str, *,
 434                          prefix=''):
 435     """
 436     Move all constraints with the specified name prefix from one bone to another.
 437     """
 438
 439     if isinstance(source, str):
 440         source = obj.pose.bones[source]
 441     if isinstance(target, str):
 442         target = obj.pose.bones[target]
 443
 444     for con in list(source.constraints):
 445         if con.name.startswith(prefix):
 446             move_constraint(source, target, con)
 447
 448
 449 ##############################################
 450 # Custom property management
 451 ##############################################
 452
 453 def deactivate_custom_properties(obj: bpy_struct, *, reset=True):
 454     """Disable drivers on custom properties and reset values to default."""
 455
 456     prefix = '["'
 457
 458     if obj != obj.id_data:
 459         prefix = obj.path_from_id() + prefix
 460
 461     adt = obj.id_data.animation_data
 462     if adt:
 463         for fcu in adt.drivers:
 464             if fcu.data_path.startswith(prefix):
 465                 fcu.mute = True
 466
 467     if reset:
 468         for key, value in obj.items():
 469             val_type = type(value)
 470             if val_type in {int, float}:
 471                 ui_data = obj.id_properties_ui(key)
 472                 rna_data = ui_data.as_dict()
 473                 obj[key] = val_type(rna_data.get("default", 0))
 474
 475
 476 def reactivate_custom_properties(obj: bpy_struct):
 477     """Re-enable drivers on custom properties."""
 478
 479     prefix = '["'
 480
 481     if obj != obj.id_data:
 482         prefix = obj.path_from_id() + prefix
 483
 484     adt = obj.id_data.animation_data
 485     if adt:
 486         for fcu in adt.drivers:
 487             if fcu.data_path.startswith(prefix):
 488                 fcu.mute = False
 489
 490
 491 def copy_custom_properties(src, dest, *, prefix='', dest_prefix='',
 492                            link_driver=False, overridable=True) -> list[tuple[str, str, Any]]:
 493     """Copy custom properties with filtering by prefix. Optionally link using drivers."""
 494     res = []
 495
 496     # Exclude addon-defined properties.
 497     exclude = {prop.identifier for prop in src.bl_rna.properties if prop.is_runtime}
 498
 499     for key, value in src.items():
 500         if key.startswith(prefix) and key not in exclude:
 501             new_key = dest_prefix + key[len(prefix):]
 502
 503             try:
 504                 ui_data_src = src.id_properties_ui(key)
 505             except TypeError:
 506                 # Some property types, e.g. Python dictionaries
 507                 # don't support id_properties_ui.
 508                 continue
 509
 510             if src != dest or new_key != key:
 511                 dest[new_key] = value
 512
 513                 dest.id_properties_ui(new_key).update_from(ui_data_src)
 514
 515                 if link_driver:
 516                     make_driver(src, quote_property(key), variables=[(dest.id_data, dest, new_key)])
 517
 518             if overridable:
 519                 dest.property_overridable_library_set(quote_property(new_key), True)
 520
 521             res.append((key, new_key, value))
 522
 523     return res
 524
 525
 526 def copy_custom_properties_with_ui(rig: 'BaseRig', src, dest_bone, *, ui_controls=None, **options):
 527     """Copy custom properties, and create rig UI for them."""
 528     if isinstance(src, str):
 529         src = rig.get_bone(src)
 530
 531     bone: PoseBone = rig.get_bone(dest_bone)
 532     mapping = copy_custom_properties(src, bone, **options)
 533
 534     if mapping:
 535         panel = rig.script.panel_with_selected_check(rig, ui_controls or rig.bones.flatten('ctrl'))
 536
 537         for key, new_key, value in sorted(mapping, key=lambda item: item[1]):
 538             name = new_key
 539
 540             # Replace delimiters with spaces
 541             if ' ' not in name:
 542                 name = re.sub(r'[_.-]', ' ', name)
 543             # Split CamelCase
 544             if ' ' not in name:
 545                 name = re.sub(r'([a-z])([A-Z])', r'\1 \2', name)
 546             # Capitalize
 547             if name.lower() == name:
 548                 name = name.title()
 549
 550             info = bone.id_properties_ui(new_key).as_dict()
 551             slider = type(value) is float and info and info.get("min", None) == 0 and info.get("max", None) == 1
 552
 553             panel.custom_prop(dest_bone, new_key, text=name, slider=slider)
 554
 555     return mapping
 556
 557
 558 ##############################################
 559 # Driver management
 560 ##############################################
 561
 562 def refresh_drivers(obj):
 563     """Cause all drivers belonging to the object to be re-evaluated, clearing any errors."""
 564
 565     # Refresh object's own drivers if any
 566     anim_data: Optional[AnimData] = getattr(obj, 'animation_data', None)
 567
 568     if anim_data:
 569         for fcu in anim_data.drivers:
 570             # Make a fake change to the driver
 571             fcu.driver.type = fcu.driver.type
 572
 573     # Material node trees aren't in any lists
 574     if isinstance(obj, Material):
 575         refresh_drivers(obj.node_tree)
 576
 577
 578 def refresh_all_drivers():
 579     """Cause all drivers in the file to be re-evaluated, clearing any errors."""
 580
 581     # Iterate over all data blocks in the file
 582     for attr in dir(bpy.data):
 583         coll = getattr(bpy.data, attr, None)
 584
 585         if isinstance(coll, bpy_prop_collection):
 586             for item in coll:
 587                 refresh_drivers(item)
 588
 589
 590 ##############################################
 591 # Utility mixin
 592 ##############################################
 593
 594 class MechanismUtilityMixin(object):
 595     obj: ArmatureObject
 596
 597     """
 598     Provides methods for more convenient creation of constraints, properties
 599     and drivers within an armature (by implicitly providing context).
 600
 601     Requires self.obj to be the armature object being worked on.
 602     """
 603
 604     def make_constraint(self, bone: str, con_type: str,
 605                         subtarget: OptionalLazy[str] = None, *,
 606                         insert_index: Optional[int] = None,
 607                         space: Optional[str] = None,
 608                         track_axis: Optional[str] = None,
 609                         use_xyz: Optional[Sequence[bool]] = None,
 610                         use_limit_xyz: Optional[Sequence[bool]] = None,
 611                         invert_xyz: Optional[Sequence[bool]] = None,
 612                         targets: Optional[list[Lazy[str | tuple | dict]]] = None,
 613                         **args):
 614         assert(self.obj.mode == 'OBJECT')
 615         return make_constraint(
 616             self.obj.pose.bones[bone], con_type, self.obj, subtarget,
 617             insert_index=insert_index, space=space, track_axis=track_axis,
 618             use_xyz=use_xyz, use_limit_xyz=use_limit_xyz, invert_xyz=invert_xyz,
 619             targets=targets,
 620             **args)
 621
 622     # noinspection PyShadowingBuiltins
 623     def make_property(self, bone: str, name: str, default, *,
 624                       min: float = 0, max: float = 1, soft_min=None, soft_max=None,
 625                       description: Optional[str] = None, overridable=True,
 626                       subtype: Optional[str] = None,
 627                       **args):
 628         assert(self.obj.mode == 'OBJECT')
 629         return make_property(
 630             self.obj.pose.bones[bone], name, default,
 631             min=min, max=max, soft_min=soft_min, soft_max=soft_max,
 632             description=description, overridable=overridable, subtype=subtype,
 633             **args
 634         )
 635
 636     # noinspection PyShadowingBuiltins
 637     def make_driver(self, owner: str | bpy_struct, prop: str,
 638                     index=-1, type='SUM',
 639                     expression: Optional[str] = None,
 640                     variables: Iterable | dict = (),
 641                     polynomial: Optional[list[float]] = None):
 642         assert(self.obj.mode == 'OBJECT')
 643         if isinstance(owner, str):
 644             owner = self.obj.pose.bones[owner]
 645         return make_driver(
 646             owner, prop, target_id=self.obj,
 647             index=index, type=type, expression=expression,
 648             variables=variables, polynomial=polynomial,
 649         )