gizmo_test3.py

import math
import os
from functools import partial
from typing import *

import bpy
from bpy import context
from bpy.utils.toolsystem import ToolDef
from bpy.types import (
	Context,
	Event,
	Gizmo,
	GizmoGroup,
	Object,
	Operator,
	Panel,
	WorkSpaceTool,
)
from bpy.props import *
import bl_ui
from bl_ui.space_toolsystem_common import activate_by_id as activate_tool
from bl_ui.space_toolsystem_toolbar import VIEW3D_PT_tools_active as view3d_tools

import bgl
from gpu_extras.presets import draw_circle_2d

from mathutils import *

icon_path = os.path.abspath("/".join((os.path.dirname(bpy.data.filepath), "icon.png")))


#==========================================================================================
def compose_world_matrix():
	GIZMO_USD_Edit.matrix_world = GIZMO_USD_Edit.matrix_basis @ GIZMO_USD_Edit.matrix_offset
	return GIZMO_USD_Edit.matrix_world


def decompose_world_matrix():
	GIZMO_USD_Edit.matrix_offset = GIZMO_USD_Edit.matrix_world @ GIZMO_USD_Edit.matrix_basis_inverse


def get_x_cb():
	return GIZMO_USD_Edit.matrix_offset.translation[0]


def set_x_cb(value):
	# print(f"Value: {value}")
	GIZMO_USD_Edit.matrix_offset.translation[0] = value
	context.active_object.matrix_world = compose_world_matrix()


def get_y_cb():
	return GIZMO_USD_Edit.matrix_offset.translation[1]


def set_y_cb(value):
	# print(f"Value: {value}")
	GIZMO_USD_Edit.matrix_offset.translation[1] = value
	context.active_object.matrix_world = compose_world_matrix()


def get_z_cb():
	return GIZMO_USD_Edit.matrix_offset.translation[2]


def set_z_cb(value):
	# print(f"Value: {value}")
	GIZMO_USD_Edit.matrix_offset.translation[2] = value
	context.active_object.matrix_world = compose_world_matrix()


#==========================================================================================OBJECT
def tag_redraw():
	for area in context.screen.areas:
		area.tag_redraw()


def get_view3d_area() -> Optional[bpy.types.SpaceView3D]:
	for area in context.screen.areas:
		for space in area.spaces:
			if space.type == "VIEW_3D":
				return space
	return None


#==========================================================================================
def get_active_tool() -> str:
	tools = context.workspace.tools
	try:
		active_tool = tools.from_space_view3d_mode("OBJECT")
		return active_tool.idname
	except AttributeError:
		return None


#==========================================================================================
class GIZMO_USD_Edit(GizmoGroup):
	bl_idname = "usd.edit"
	bl_label = "USD Edit"
	bl_space_type = 'VIEW_3D'
	bl_region_type = 'WINDOW'
	bl_options = {
		'3D',
		'PERSISTENT',
		'SHOW_MODAL_ALL'
	}
	
	matrix_basis         = Matrix.Identity(4)
	matrix_basis_inverse = Matrix.Identity(4)
	matrix_offset        = Matrix.Identity(4)
	matrix_world         = Matrix.Identity(4)

	active_axis = None
	needs_reset = False

	@classmethod
	def poll(cls, context):
		active_tool = get_active_tool()
		ob = context.active_object
		result = all([active_tool == TOOL_GGT_gizmo_object.bl_idname, ob])
		if not result:
			GIZMO_USD_Edit.needs_reset = True
		return result

	def invoke_prepare(self, context, gizmo):
		print("Invoke prepare")

	def setup(self, context):
		# print("setup")
		self.alpha = 1.0
		self.x_color = (1.0, 0.09, 0.267)
		self.y_axis_color = (0.545, 0.863, 0.0)
		self.z_axis_color = (0.161, 0.384, 1.0)
		self.select_color = (0.812, 0.812, 0.812)
		self.gizmo_scale = 1.0
		self.last_draw_moving = False

		self.init_matrices(context.active_object.matrix_world)

		self.arrow_x = self.gizmos.new("GIZMO_GT_arrow_3d")		
		self.arrow_y = self.gizmos.new("GIZMO_GT_arrow_3d")
		self.arrow_z = self.gizmos.new("GIZMO_GT_arrow_3d")

		self.arrow_x.target_set_handler("offset", get=get_x_cb, set=set_x_cb)
		self.arrow_y.target_set_handler("offset", get=get_y_cb, set=set_y_cb)
		self.arrow_z.target_set_handler("offset", get=get_z_cb, set=set_z_cb)

		self.init_handles()

		self.last_modal = [False, False, False]

	def gizmos_modal_get(self, invert=False):
		if invert:
			result = [not x.is_modal for x in self.gizmos]
		else:
			result = [x.is_modal for x in self.gizmos]
		return result

	def is_moving(self):
		return any(self.gizmos_modal_get())

	def init_matrices(self, basis:Matrix): 
		GIZMO_USD_Edit.matrix_basis = basis.copy()
		GIZMO_USD_Edit.matrix_basis_inverse = basis.inverted_safe()
		GIZMO_USD_Edit.matrix_offset = Matrix.Identity(4)

		## initial world is same as basis-- offset is zero
		GIZMO_USD_Edit.matrix_world = basis.copy()
		
		self.x_matrix_offset = Matrix.Rotation(math.pi * 0.5, 4, "Y")
		self.y_matrix_offset = Matrix.Rotation(math.pi * -0.5, 4, "X")
		self.z_matrix_offset = Matrix.Identity(4)

		self.x_matrix = (GIZMO_USD_Edit.matrix_basis @ self.x_matrix_offset).normalized()
		self.y_matrix = (GIZMO_USD_Edit.matrix_basis @ self.y_matrix_offset).normalized()
		self.z_matrix = GIZMO_USD_Edit.matrix_basis.copy()

	def refresh(self, context:Context):
		print("refresh")
		# self.init_matrices(context.active_object.matrix_world)
		if not sum(self.last_modal) == 0:
			return

		# self.init_handles()
		# print(self.last_modal)
		# for modal, gizmo in zip(self.last_modal, self.gizmos):
		# 	if not modal:
		# 		gizmo.matrix_basis = gizmo.matrix_basis @ GIZMO_USD_Edit.matrix_offset

	def init_handles(self):
		self.arrow_x.alpha = self.alpha
		self.arrow_x.color = self.x_color
		self.arrow_x.color_highlight = self.select_color
		self.arrow_x.scale_basis = self.gizmo_scale
		self.arrow_x.matrix_basis = self.x_matrix.copy()

		self.arrow_y.alpha = self.alpha
		self.arrow_y.color = self.y_axis_color
		self.arrow_y.color_highlight = self.select_color
		self.arrow_y.scale_basis = self.gizmo_scale
		self.arrow_y.matrix_basis = self.y_matrix.copy()

		self.arrow_z.alpha = self.alpha
		self.arrow_z.color = self.z_axis_color
		self.arrow_z.color_highlight = self.select_color
		self.arrow_z.scale_basis = self.gizmo_scale
		self.arrow_z.matrix_basis = self.z_matrix.copy()

		for giz in self.gizmos:
			giz.matrix_offset = Matrix.Identity(4)

	def draw_prepare(self, context:Context):
		gizmo_scale = 2.0
		alpha = 1.0
		x_color = (1,0,0)
		y_axis_color = (0,1,0)
		z_axis_color = (0,0,1)
		select_color = (0.82, 0.82, 0.82)

		moving = self.is_moving()
		visibility = [False, False, False]

		if moving and not self.last_draw_moving:
			print("started move")
			# self.init_matrices(context.active_object.matrix_world)
			# self.init_handles()
			visibility = self.gizmos_modal_get(invert=True)
		elif moving and self.last_draw_moving:
			pass
			visibility = self.gizmos_modal_get(invert=True)
		elif not moving and self.last_draw_moving:
			print("finishing move")
			## this communicates out to the other handler
			visibility = [True, True, True]
			GIZMO_USD_Edit.needs_reset = True

		if GIZMO_USD_Edit.needs_reset:
			self.init_matrices(context.active_object.matrix_world)
			self.init_handles()
			# compose_world_matrix()
			# for gizmo in self.gizmos:
			# 	gizmo.target_set_value("offset", [0.0])
			GIZMO_USD_Edit.needs_reset = False
			# tag_redraw()

		self.last_draw_moving = moving
		self.last_modal = self.gizmos_modal_get()

		## uncomment this and the interaction breaks
		# for vis, giz in zip(visibility, self.gizmos):
		# 	giz.hide = vis


#==========================================================================================OBJECT
class TOOL_GGT_gizmo_object(WorkSpaceTool):
	bl_space_type = 'VIEW_3D'
	bl_context_mode = 'OBJECT'
	bl_idname = "tool.usd_edit"
	bl_label = "Usd Edit"
	bl_icon = icon_path

	bl_widget = "usd.edit"
	bl_cursor = 'PAINT_CROSS' #cursor()

	@classmethod
	def poll(cls, context:Context) -> bool:
		return context.active_object and context.active_object.type == "USDPRIM"

	def setup(self, idname):
		print("TOOL_GGT_gizmo_object setup")
		print(idname)

	def draw_settings(context, layout, tool):
		layout.label(text="Settings area")


#==========================================================================================OBJECT
classes = [
	GIZMO_USD_Edit,
]

tools = [
	TOOL_GGT_gizmo_object,
]


def register():
	print("Register")

	for cls in classes:
		bpy.utils.register_class(cls)

	for tool in tools:
		separator = tool is tools[0]
		bpy.utils.register_tool(tool, separator=separator, group=False)


def unregister():
	for tool in reversed(tools):
		try:
			bpy.utils.unregister_tool(tool)
		except (AttributeError, RuntimeError):
			continue

	for cls in reversed(classes):
		try:
			bpy.utils.unregister_class(cls)
		except (AttributeError, RuntimeError):
			continue


if __name__ == "__main__":
	unregister()
	register()