Files
2026-03-17 14:58:51 -06:00

223 lines
7.5 KiB
Python

# SPDX-FileCopyrightText: 2025 Blender Studio Tools Authors
#
# SPDX-License-Identifier: GPL-3.0-or-later
import bpy
from bpy.types import GizmoGroup, Gizmo, PoseBone, Operator
from .utils import get_addon_prefs
### MSGBUS FUNCTIONS ###
# The Gizmo API doesn't provide the necessary callbacks to do careful partial
# updates of gizmo data based on what properties are modified.
# However, the msgbus system allows us to create our own connections between
# Blender properties and the functions that should be called when those properties
# change.
# msgbus system needs an arbitrary python object as storage, so here it is.
gizmo_msgbus = object()
def mb_ensure_gizmos_on_active_armature(gizmo_group):
"""Ensure gizmos exist for all PoseBones that need them."""
context = bpy.context
obj = context.object
for pose_bone in obj.pose.bones:
try:
if pose_bone.enable_bone_gizmo and pose_bone.name not in gizmo_group.widgets:
gizmo = gizmo_group.create_gizmo(context, pose_bone)
elif not pose_bone.bone_gizmo.shape_object:
for bone_name, gizmo in gizmo_group.widgets.items():
if bone_name == pose_bone.name:
gizmo.custom_shape = None
except ReferenceError:
# StructRNA of type BoneGizmoGroup has been removed.
# TODO: Not sure when this happens.
pass
def mb_refresh_all_gizmo_colors(gizmo_group):
"""Keep Gizmo colors in sync with addon preferences."""
context = bpy.context
try:
for bone_name, gizmo in gizmo_group.widgets.items():
gizmo.refresh_colors(context)
except ReferenceError:
# StructRNA of type BoneGizmoGroup has been removed.
# TODO: Not sure when this happens.
pass
def mb_refresh_single_gizmo(gizmo_group, bone_name):
"""Refresh Gizmo behaviour settings. This should be called when the user changes
the Gizmo settings in the Properties editor.
"""
if not gizmo_group:
return
context = bpy.context
pose_bone = context.object.pose.bones.get(bone_name)
gizmo_props = pose_bone.bone_gizmo
try:
gizmo = gizmo_group.widgets[bone_name]
except:
return
if gizmo_props.operator != 'None':
op_name = gizmo_props.operator
if op_name == 'transform.rotate' and gizmo_props.rotation_mode == 'TRACKBALL':
op_name = 'transform.trackball'
op = gizmo.target_set_operator(op_name)
if op_name == 'transform.rotate' and gizmo_props.rotation_mode in 'XYZ':
op.orient_type = 'LOCAL'
op.constraint_axis = [axis == gizmo_props.rotation_mode for axis in 'XYZ']
if op_name in ['transform.translate', 'transform.resize']:
op.constraint_axis = gizmo_props.transform_axes
gizmo.init_properties(context)
def mb_refresh_single_gizmo_shape(gizmo_group, bone_name):
"""Re-calculate a gizmo's vertex indicies. This is expensive, so it should
be called sparingly."""
context = bpy.context
print("refresh single gizmo shape")
if not context.object or context.object.type != 'ARMATURE' or context.object.mode != 'POSE':
return
try:
gizmo = gizmo_group.widgets[bone_name]
gizmo.init_shape(context)
except:
pass
class BoneGizmoGroup(GizmoGroup):
"""This single GizmoGroup manages all bone gizmos for all rigs.""" # TODO: Currently this will have issues when there are two rigs with similar bone names. Rig object names should be included when identifying widgets.
bl_idname = "OBJECT_GGT_bone_gizmo"
bl_label = "Bone Gizmos"
bl_space_type = 'VIEW_3D'
bl_region_type = 'WINDOW'
bl_options = {
'3D' # Lets Gizmos use the 'draw_select' function to draw into a selection pass.
,'PERSISTENT'
,'SHOW_MODAL_ALL' # TODO what is this
,'DEPTH_3D' # Provides occlusion but results in Z-fighting when gizmo geometry isn't offset from the source mesh.
,'SELECT' # I thought this would make Gizmo.select do something but doesn't seem that way
,'SCALE' # This makes all gizmos' scale relative to the world rather than the camera, so we don't need to set use_draw_scale on each Gizmo. (And that option does nothing because of this one)
}
@classmethod
def poll(cls, context):
return context.scene.bone_gizmos_enabled and context.object \
and context.object.type == 'ARMATURE' and context.object.mode=='POSE'
def setup(self, context):
"""Runs when this GizmoGroup is created, I think.
Executed by Blender on launch, and for some reason
also when changing bone group colors. WHAT!?"""
self.widgets = {}
mb_ensure_gizmos_on_active_armature(self)
# Hook up the addon preferences to color refresh function
# using msgbus system.
addon_prefs = get_addon_prefs(context)
global gizmo_msgbus
bpy.msgbus.subscribe_rna(
key = addon_prefs.path_resolve('preferences', False)
,owner = gizmo_msgbus
,args = (self,)
,notify = mb_refresh_all_gizmo_colors
)
# Hook up Custom Gizmo checkbox to a function that will ensure that
# a Gizmo instance actually exists for each bone that needs one.
bpy.msgbus.subscribe_rna(
key = (PoseBone, "enable_bone_gizmo")
,owner = gizmo_msgbus
,args = (self,)
,notify = mb_ensure_gizmos_on_active_armature
)
def create_gizmo(self, context, pose_bone) -> Gizmo:
"""Add a gizmo to this GizmoGroup based on user-defined properties."""
gizmo_props = pose_bone.bone_gizmo
if not pose_bone.enable_bone_gizmo:
return
gizmo = self.gizmos.new('GIZMO_GT_bone_gizmo')
gizmo.bone_name = pose_bone.name
# Hook up gizmo properties (the ones that can be customized by user)
# to the gizmo refresh functions, using msgbus system.
global gizmo_msgbus
bpy.msgbus.subscribe_rna(
key = gizmo_props
,owner = gizmo_msgbus
,args = (self, gizmo.bone_name)
,notify = mb_refresh_single_gizmo
)
for prop_name in ['shape_object', 'vertex_group_name', 'face_map_name', 'use_face_map']:
bpy.msgbus.subscribe_rna(
key = gizmo_props.path_resolve(prop_name, False)
,owner = gizmo_msgbus
,args = (self, gizmo.bone_name)
,notify = mb_refresh_single_gizmo_shape
)
self.widgets[pose_bone.name] = gizmo
mb_refresh_single_gizmo(self, pose_bone.name)
mb_refresh_single_gizmo_shape(self, pose_bone.name)
return gizmo
def refresh(self, context):
"""This is a Gizmo API function, called by Blender on what seems to be
depsgraph updates and frame changes.
Refresh all visible gizmos that use vertex group masking.
This should be done whenever a bone position changes.
This should be kept performant!
"""
dg = context.evaluated_depsgraph_get()
eval_meshes = {}
for bonename, gizmo in self.widgets.items():
pb = gizmo.get_pose_bone(context)
if not gizmo or not gizmo.is_using_vgroup(context) or not gizmo.poll(context):
continue
obj = pb.bone_gizmo.shape_object
if obj.name in eval_meshes:
eval_mesh = eval_meshes[obj.name]
else:
eval_meshes[obj.name] = eval_mesh = obj.evaluated_get(dg).to_mesh()
eval_mesh.calc_loop_triangles()
try:
gizmo.refresh_shape_vgroup(context, eval_mesh)
except ReferenceError:
# For some reason sometimes it complains that StructRNA of the Mesh has been removed. I don't get why.
pass
class BONEGIZMO_OT_RestartGizmoGroup(Operator):
"""Re-initialize all gizmos. Needed when the gizmo shape objects are modified, since there's no dependency between gizmos and their target shapes"""
bl_idname = "pose.restart_gizmos"
bl_label = "Refresh Gizmos"
bl_options = {'REGISTER', 'UNDO'}
def execute(self, context):
bpy.utils.unregister_class(BoneGizmoGroup)
bpy.utils.register_class(BoneGizmoGroup)
return {'FINISHED'}
registry = [
BoneGizmoGroup,
BONEGIZMO_OT_RestartGizmoGroup
]
def unregister():
# Unhook everything from msgbus system
global gizmo_msgbus
bpy.msgbus.clear_by_owner(gizmo_msgbus)