rigify/utils/objects.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6
   7 from typing import TYPE_CHECKING
   8 from bpy.types import LayerCollection, Collection, Object
   9
  10 from .misc import ArmatureObject
  11 from .naming import strip_org
  12
  13 from mathutils import Matrix
  14
  15 if TYPE_CHECKING:
  16     from ..generate import Generator
  17     from ..base_rig import BaseRig
  18
  19
  20 # noinspection SpellCheckingInspection
  21 def create_object_data(obj_type, name):
  22     if obj_type == 'EMPTY':
  23         return None
  24     if obj_type == 'MESH':
  25         return bpy.data.meshes.new(name)
  26     if obj_type in ('CURVE', 'SURFACE', 'FONT'):
  27         return bpy.data.curves.new(name, obj_type)
  28     if obj_type == 'META':
  29         return bpy.data.metaballs.new(name)
  30     if obj_type == 'CURVES':
  31         return bpy.data.hair_curves.new(name)
  32     if obj_type == 'POINTCLOUD':
  33         return bpy.data.pointclouds.new(name)
  34     if obj_type == 'VOLUME':
  35         return bpy.data.volumes.new(name)
  36     if obj_type == 'GREASEPENCIL':
  37         return bpy.data.grease_pencils.new(name)
  38     if obj_type == 'ARMATURE':
  39         return bpy.data.armatures.new(name)
  40     if obj_type == 'LATTICE':
  41         return bpy.data.lattices.new(name)
  42     raise ValueError(f"Invalid object type {obj_type}")
  43
  44
  45 class ArtifactManager:
  46     generator: 'Generator'
  47
  48     collection: Collection | None
  49     layer_collection: LayerCollection | None
  50
  51     used_artifacts: list[Object]
  52     temp_artifacts: list[Object]
  53
  54     artifact_reuse_table: dict[tuple[str, ...], Object]
  55
  56     def __init__(self, generator: 'Generator'):
  57         self.generator = generator
  58         self.collection = None
  59         self.layer_collection = None
  60         self.used_artifacts = []
  61         self.temp_artifacts = []
  62         self.artifact_reuse_table = {}
  63
  64     def _make_name(self, owner: 'BaseRig', name: str):
  65         return self.generator.obj.name + ":" + strip_org(owner.base_bone) + ":" + name
  66
  67     def create_new(self, owner: 'BaseRig', obj_type: str, name: str):
  68         """
  69         Creates an artifact object of the specified type and name. If it already exists, all
  70         references are updated to point to the new instance, and the existing one is deleted.
  71
  72         Parameters:
  73             owner: rig component that requests the object.
  74             obj_type: type of the object to create.
  75             name: unique name of the object within the rig component.
  76         Returns:
  77             Object that was created.
  78         """
  79         return self.find_or_create(owner, obj_type, name, recreate=True)[1]
  80
  81     def find_or_create(self, owner: 'BaseRig', obj_type: str, name: str, *, recreate=False):
  82         """
  83         Creates or reuses an artifact object of the specified type.
  84
  85         Parameters:
  86             owner: rig component that requests the object.
  87             obj_type: type of the object to create.
  88             name: unique name of the object within the rig component.
  89             recreate: instructs that the object should be re-created from scratch even if it exists.
  90         Returns:
  91             (bool, Object) tuple, with the boolean specifying if the object already existed.
  92         """
  93
  94         obj_name = self._make_name(owner, name)
  95         key = (owner.base_bone, name)
  96
  97         obj = self.artifact_reuse_table.get(key)
  98
  99         # If the existing object has incorrect type, delete it
 100         if obj and obj.type != obj_type:
 101             if obj in self.used_artifacts:
 102                 owner.raise_error(f"duplicate reuse of artifact object {obj.name}")
 103
 104             print(f"RIGIFY: incompatible artifact object {obj.name} type: {obj.type} instead of {obj_type}")
 105             del self.artifact_reuse_table[key]
 106             bpy.data.objects.remove(obj)
 107             obj = None
 108
 109         # Reuse the existing object
 110         if obj:
 111             if obj in self.used_artifacts:
 112                 owner.raise_error(f"duplicate reuse of artifact object {obj.name}")
 113
 114             if recreate:
 115                 # Forcefully re-create and replace the existing object
 116                 obj.name += '-OLD'
 117                 if data := obj.data:
 118                     data.name += '-OLD'
 119
 120                 new_obj = bpy.data.objects.new(obj_name, create_object_data(obj_type, obj_name))
 121
 122                 obj.user_remap(new_obj)
 123                 self.artifact_reuse_table[key] = new_obj
 124                 bpy.data.objects.remove(obj)
 125                 obj = new_obj
 126
 127             # Ensure the existing object is visible
 128             obj.hide_viewport = False
 129             obj.hide_set(False, view_layer=self.generator.view_layer)
 130
 131             if not obj.visible_get(view_layer=self.generator.view_layer):
 132                 owner.raise_error(f"could not un-hide existing artifact object {obj.name}")
 133
 134             # Try renaming the existing object
 135             obj.name = obj_name
 136             if data := obj.data:
 137                 data.name = obj_name
 138
 139             found = True
 140
 141         # Create an object from scratch
 142         else:
 143             obj = bpy.data.objects.new(obj_name, create_object_data(obj_type, obj_name))
 144
 145             self.generator.collection.objects.link(obj)
 146             self.artifact_reuse_table[key] = obj
 147
 148             found = False
 149
 150         self.used_artifacts.append(obj)
 151
 152         obj.rigify_owner_rig = self.generator.obj
 153         obj["rigify_artifact_id"] = key
 154
 155         obj.parent = self.generator.obj
 156         obj.parent_type = 'OBJECT'
 157         obj.matrix_parent_inverse = Matrix.Identity(4)
 158         obj.matrix_basis = Matrix.Identity(4)
 159
 160         return found, obj
 161
 162     def new_temporary(self, owner: 'BaseRig', obj_type: str, name="temp"):
 163         """
 164         Creates a new temporary object of the specified type.
 165         The object will be removed after generation finishes.
 166         """
 167         obj_name = "TEMP:" + self._make_name(owner, name)
 168         obj = bpy.data.objects.new(obj_name, create_object_data(obj_type, obj_name))
 169         obj.rigify_owner_rig = self.generator.obj
 170         obj["rigify_artifact_id"] = 'temporary'
 171         self.generator.collection.objects.link(obj)
 172         self.temp_artifacts.append(obj)
 173         return obj
 174
 175     def remove_temporary(self, obj):
 176         """
 177         Immediately removes a temporary object previously created using new_temporary.
 178         """
 179         self.temp_artifacts.remove(obj)
 180         bpy.data.objects.remove(obj)
 181
 182     def generate_init_existing(self, armature: ArmatureObject):
 183         for obj in bpy.data.objects:
 184             if obj.rigify_owner_rig != armature:
 185                 continue
 186
 187             aid = obj["rigify_artifact_id"]
 188             if isinstance(aid, list) and all(isinstance(x, str) for x in aid):
 189                 self.artifact_reuse_table[tuple(aid)] = obj
 190             else:
 191                 print(f"RIGIFY: removing orphan artifact {obj.name}")
 192                 bpy.data.objects.remove(obj)
 193
 194     def generate_cleanup(self):
 195         for obj in self.temp_artifacts:
 196             bpy.data.objects.remove(obj)
 197
 198         self.temp_artifacts = []
 199
 200         for key, obj in self.artifact_reuse_table.items():
 201             if obj in self.used_artifacts:
 202                 obj.hide_viewport = True
 203                 obj.hide_render = True
 204             else:
 205                 del self.artifact_reuse_table[key]
 206                 bpy.data.objects.remove(obj)