rigify/utils/naming.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import random
   6 import time
   7 import re
   8 import collections
   9 import enum
  10
  11 from typing import Optional, TYPE_CHECKING
  12
  13
  14 if TYPE_CHECKING:
  15     from ..base_generate import BaseGenerator
  16
  17
  18 ORG_PREFIX = "ORG-"  # Prefix of original bones.
  19 MCH_PREFIX = "MCH-"  # Prefix of mechanism bones.
  20 DEF_PREFIX = "DEF-"  # Prefix of deformation bones.
  21 ROOT_NAME = "root"   # Name of the root bone.
  22
  23 _PREFIX_TABLE = {'org': "ORG", 'mch': "MCH", 'def': "DEF", 'ctrl': ''}
  24
  25 ########################################################################
  26 # Name structure
  27 ########################################################################
  28
  29 NameParts = collections.namedtuple('NameParts', ['prefix', 'base', 'side_z', 'side', 'number'])
  30
  31
  32 def split_name(name: str):
  33     name_parts = re.match(
  34         r'^(?:(ORG|MCH|DEF)-)?(.*?)([._-][tTbB])?([._-][lLrR])?(?:\.(\d+))?$', name)
  35     return NameParts(*name_parts.groups())
  36
  37
  38 def is_control_bone(name: str):
  39     return not split_name(name).prefix
  40
  41
  42 def combine_name(parts: NameParts, *, prefix=None, base=None, side_z=None, side=None, number=None):
  43     eff_prefix = prefix if prefix is not None else parts.prefix
  44     eff_number = number if number is not None else parts.number
  45     if isinstance(eff_number, int):
  46         eff_number = '%03d' % eff_number
  47
  48     return ''.join([
  49         eff_prefix+'-' if eff_prefix else '',
  50         base if base is not None else parts.base,
  51         side_z if side_z is not None else parts.side_z or '',
  52         side if side is not None else parts.side or '',
  53         '.'+eff_number if eff_number else '',
  54     ])
  55
  56
  57 def insert_before_lr(name: str, text: str) -> str:
  58     parts = split_name(name)
  59
  60     if parts.side:
  61         return combine_name(parts, base=parts.base + text)
  62     else:
  63         return name + text
  64
  65
  66 def make_derived_name(name: str, subtype: str, suffix: Optional[str] = None):
  67     """ Replaces the name prefix, and optionally adds the suffix (before .LR if found).
  68     """
  69     assert(subtype in _PREFIX_TABLE)
  70
  71     parts = split_name(name)
  72     new_base = parts.base + (suffix or '')
  73
  74     return combine_name(parts, prefix=_PREFIX_TABLE[subtype], base=new_base)
  75
  76
  77 ########################################################################
  78 # Name mirroring
  79 ########################################################################
  80
  81 class Side(enum.IntEnum):
  82     LEFT = -1
  83     MIDDLE = 0
  84     RIGHT = 1
  85
  86     @staticmethod
  87     def from_parts(parts: NameParts):
  88         if parts.side:
  89             if parts.side[1].lower() == 'l':
  90                 return Side.LEFT
  91             else:
  92                 return Side.RIGHT
  93         else:
  94             return Side.MIDDLE
  95
  96     @staticmethod
  97     def to_string(parts: NameParts, side: 'Side'):
  98         if side != Side.MIDDLE:
  99             side_char = 'L' if side == Side.LEFT else 'R'
 100             side_str = parts.side or parts.side_z
 101
 102             if side_str:
 103                 sep, side_char2 = side_str[0:2]
 104                 if side_char2.lower() == side_char2:
 105                     side_char = side_char.lower()
 106             else:
 107                 sep = '.'
 108
 109             return sep + side_char
 110         else:
 111             return ''
 112
 113     @staticmethod
 114     def to_name(parts: NameParts, side: 'Side'):
 115         new_side = Side.to_string(parts, side)
 116         return combine_name(parts, side=new_side)
 117
 118
 119 class SideZ(enum.IntEnum):
 120     TOP = 2
 121     MIDDLE = 0
 122     BOTTOM = -2
 123
 124     @staticmethod
 125     def from_parts(parts: NameParts):
 126         if parts.side_z:
 127             if parts.side_z[1].lower() == 't':
 128                 return SideZ.TOP
 129             else:
 130                 return SideZ.BOTTOM
 131         else:
 132             return SideZ.MIDDLE
 133
 134     @staticmethod
 135     def to_string(parts: NameParts, side: 'SideZ'):
 136         if side != SideZ.MIDDLE:
 137             side_char = 'T' if side == SideZ.TOP else 'B'
 138             side_str = parts.side_z or parts.side
 139
 140             if side_str:
 141                 sep, side_char2 = side_str[0:2]
 142                 if side_char2.lower() == side_char2:
 143                     side_char = side_char.lower()
 144             else:
 145                 sep = '.'
 146
 147             return sep + side_char
 148         else:
 149             return ''
 150
 151     @staticmethod
 152     def to_name(parts: NameParts, side: 'SideZ'):
 153         new_side = SideZ.to_string(parts, side)
 154         return combine_name(parts, side_z=new_side)
 155
 156
 157 NameSides = collections.namedtuple('NameSides', ['base', 'side', 'side_z'])
 158
 159
 160 def get_name_side(name: str):
 161     return Side.from_parts(split_name(name))
 162
 163
 164 def get_name_side_z(name: str):
 165     return SideZ.from_parts(split_name(name))
 166
 167
 168 def get_name_base_and_sides(name: str):
 169     parts = split_name(name)
 170     base = combine_name(parts, side='', side_z='')
 171     return NameSides(base, Side.from_parts(parts),  SideZ.from_parts(parts))
 172
 173
 174 def change_name_side(name: str,
 175                      side: Optional[Side] = None, *,
 176                      side_z: Optional[SideZ] = None):
 177     parts = split_name(name)
 178     new_side = None if side is None else Side.to_string(parts, side)
 179     new_side_z = None if side_z is None else SideZ.to_string(parts, side_z)
 180     return combine_name(parts, side=new_side, side_z=new_side_z)
 181
 182
 183 def mirror_name(name: str):
 184     parts = split_name(name)
 185     side = Side.from_parts(parts)
 186
 187     if side != Side.MIDDLE:
 188         return Side.to_name(parts, -side)
 189     else:
 190         return name
 191
 192
 193 def mirror_name_z(name: str):
 194     parts = split_name(name)
 195     side = SideZ.from_parts(parts)
 196
 197     if side != SideZ.MIDDLE:
 198         return SideZ.to_name(parts, -side)
 199     else:
 200         return name
 201
 202
 203 ########################################################################
 204 # Name manipulation
 205 ########################################################################
 206
 207 def get_name(bone) -> Optional[str]:
 208     return bone.name if bone else None
 209
 210
 211 def strip_trailing_number(name: str):
 212     return combine_name(split_name(name), number='')
 213
 214
 215 def strip_prefix(name: str):
 216     return combine_name(split_name(name), prefix='')
 217
 218
 219 def unique_name(collection, base_name: str):
 220     parts = split_name(base_name)
 221     name = combine_name(parts, number='')
 222     count = 1
 223
 224     while name in collection:
 225         name = combine_name(parts, number=count)
 226         count += 1
 227
 228     return name
 229
 230
 231 def strip_org(name: str):
 232     """ Returns the name with ORG_PREFIX stripped from it.
 233     """
 234     if name.startswith(ORG_PREFIX):
 235         return name[len(ORG_PREFIX):]
 236     else:
 237         return name
 238
 239
 240 org_name = strip_org
 241
 242
 243 def strip_mch(name: str):
 244     """ Returns the name with MCH_PREFIX stripped from it.
 245         """
 246     if name.startswith(MCH_PREFIX):
 247         return name[len(MCH_PREFIX):]
 248     else:
 249         return name
 250
 251
 252 def strip_def(name: str):
 253     """ Returns the name with DEF_PREFIX stripped from it.
 254         """
 255     if name.startswith(DEF_PREFIX):
 256         return name[len(DEF_PREFIX):]
 257     else:
 258         return name
 259
 260
 261 def org(name: str):
 262     """ Prepends the ORG_PREFIX to a name if it doesn't already have
 263         it, and returns it.
 264     """
 265     if name.startswith(ORG_PREFIX):
 266         return name
 267     else:
 268         return ORG_PREFIX + name
 269
 270
 271 make_original_name = org
 272
 273
 274 def mch(name: str):
 275     """ Prepends the MCH_PREFIX to a name if it doesn't already have
 276         it, and returns it.
 277     """
 278     if name.startswith(MCH_PREFIX):
 279         return name
 280     else:
 281         return MCH_PREFIX + name
 282
 283
 284 make_mechanism_name = mch
 285
 286
 287 def deformer(name: str):
 288     """ Prepends the DEF_PREFIX to a name if it doesn't already have
 289         it, and returns it.
 290     """
 291     if name.startswith(DEF_PREFIX):
 292         return name
 293     else:
 294         return DEF_PREFIX + name
 295
 296
 297 make_deformer_name = deformer
 298
 299
 300 def random_id(length=8):
 301     """ Generates a random alphanumeric id string.
 302     """
 303     t_length = int(length / 2)
 304     r_length = int(length / 2) + int(length % 2)
 305
 306     chars = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', 'a', 'b', 'c', 'd', 'e', 'f',
 307              'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w',
 308              'x', 'y', 'z']
 309     text = ""
 310     for i in range(0, r_length):
 311         text += random.choice(chars)
 312     text += str(hex(int(time.time())))[2:][-t_length:].rjust(t_length, '0')[::-1]
 313     return text
 314
 315
 316 def choose_derived_bone(generator: 'BaseGenerator', original: str, subtype: str, *,
 317                         by_owner=True, recursive=True):
 318     bones = generator.obj.pose.bones
 319     names = generator.find_derived_bones(original, by_owner=by_owner, recursive=recursive)
 320
 321     direct = make_derived_name(original, subtype)
 322     if direct in names and direct in bones:
 323         return direct
 324
 325     prefix = _PREFIX_TABLE[subtype] + '-'
 326     matching = [name for name in names if name.startswith(prefix) and name in bones]
 327
 328     if len(matching) > 0:
 329         return matching[0]
 330
 331     # Try matching bones created by legacy rigs just by name - there is no origin data
 332     from ..base_generate import LegacyRig
 333
 334     if isinstance(generator.bone_owners.get(direct), LegacyRig):
 335         if not by_owner or generator.bone_owners.get(original) is generator.bone_owners[direct]:
 336             assert direct in bones
 337             return direct
 338
 339     return None
 340
 341
 342 _MIRROR_MAP_RAW = [
 343     ("Left", "Right"),
 344     ("L", "R"),
 345 ]
 346 _MIRROR_MAP = {
 347     **{a: b for a, b in _MIRROR_MAP_RAW},
 348     **{b: a for a, b in _MIRROR_MAP_RAW},
 349     **{a.lower(): b.lower() for a, b in _MIRROR_MAP_RAW},
 350     **{b.lower(): a.lower() for a, b in _MIRROR_MAP_RAW},
 351 }
 352 _MIRROR_RE = [
 353     r"(?<![a-z])(left|light)(?![a-z])",
 354     r"(?<=[\._])(l|r)(?![a-z])",
 355 ]
 356
 357
 358 def mirror_name_fuzzy(name: str) -> str:
 359     """Try to mirror a name by trying various patterns without expecting any rigid structure."""
 360
 361     for reg in _MIRROR_RE:
 362         new_name = re.sub(reg, lambda m: _MIRROR_MAP.get(m[0], m[0]), name, flags=re.IGNORECASE)
 363         if new_name != name:
 364             return new_name
 365
 366     return name