__init__.py

@@ -24,6 +24,8 @@ bl_info = {
 }
 
 import bpy, importlib
+from bpy.types import AddonPreferences
+from bpy.props import BoolProperty
 
 from . import smart_weight_transfer
 from . import force_apply_mirror
@@ -40,11 +42,29 @@ modules = [
 	weight_paint_context_menu
 ]
 
+class EasyWeightPreferences(AddonPreferences):
+	bl_idname = __name__
+
+	debug_mode: BoolProperty(
+		name='Debug UI',
+		description='When enabled, add an EasyWeight panel to the Sidebar, used for debugging',
+		default=False,
+	)
+	
+	def draw(self, context):
+		layout = self.layout
+
+		layout.prop(self, 'debug_mode')
+
 def register():
+	from bpy.utils import register_class
+	register_class(EasyWeightPreferences)
 	for m in modules:
 		importlib.reload(m)
 		m.register()
 
 def unregister():
+	from bpy.utils import unregister_class
+	unregister_class(EasyWeightPreferences)
 	for m in modules:
 		m.unregister()

toggle_weight_paint.py

@@ -13,8 +13,6 @@ import bpy
 
 # When running the operator again, it should restore everything to how it was before.
 
-coll_name = "Weight Paint Toggle Helper"
-
 def get_armature_of_meshob(obj: bpy.types.Object):
 	"""Find and return the armature that deforms this mesh object."""
 	for m in obj.modifiers:
@@ -83,27 +81,26 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 		# Enter WP mode.
 		bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
 
-		# If there is an armature, make sure it is visible.
+		### ENSURING ARMATURE VISIBILITY
 		armature = get_armature_of_meshob(obj)
 		if not armature:
 			return {'FINISHED'}
-		context.screen['wpt']['armature_enabled'] = armature.hide_viewport
-		context.screen['wpt']['armature_hide'] = armature.hide_get()
-		context.screen['wpt']['armature_in_front'] = armature.show_in_front
+		# Save all object visibility related info so it can be restored later.
+		context.screen['wpt']['arm_enabled'] = armature.hide_viewport
+		context.screen['wpt']['arm_hide'] = armature.hide_get()
+		context.screen['wpt']['arm_in_front'] = armature.show_in_front
+		context.screen['wpt']['arm_coll_assigned'] = False
 		armature.hide_viewport = False
 		armature.hide_set(False)
 		armature.show_in_front = True
+		if context.space_data.local_view:
+			context.screen['wpt']['arm_local_view'] = armature.local_view_get(context.space_data)
+			armature.local_view_set(context.space_data, True)
 
-		# If the armature is still not visible, use a collection.
-		if not armature.visible_get():
-			coll = bpy.data.collections.get(coll_name)
-			if not coll:
-				coll = bpy.data.collections.new(coll_name)
-			if coll_name not in context.scene.collection.children:
-				context.scene.collection.children.link(coll)
-			
-			if armature.name not in coll.objects:
-				coll.objects.link(armature)
+		# If the armature is still not visible, add it to the scene root collection.
+		if not armature.visible_get() and not armature.name in context.scene.collection.objects:
+			context.scene.collection.objects.link(armature)
+			context.screen['wpt']['arm_coll_assigned'] = True
 
 		if not armature.visible_get():
 			self.report({'WARNING'}, "Failed to make the armature visible. Curse you, Blender!!!")
@@ -125,7 +122,7 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 			obj.display_type = obj['wpt_display_type']
 			del obj['wpt_display_type']
 
-		if 'wpt' not in context.screen:
+		if 'wpt' not in context.screen or 'mode' not in context.screen['wpt'].to_dict():
 			# There is no saved data to restore from, nothing else to do.
 			bpy.ops.object.mode_set(mode='OBJECT')
 			return {'FINISHED'}
@@ -135,26 +132,37 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 		# Restore mode.
 		bpy.ops.object.mode_set(mode=wpt['mode'])
 
+		# Reset the stored data
+		context.screen['wpt'] = {}
 		# Flag to save that the last time the operator ran we were EXITING wp mode.
 		context.screen['wpt']['last_switch_in'] = False
 
 		# Restore shading options.
+		if 'light' in wpt:	# If we stored solid view shading settings, restore those settings but don't restore the shading type.
+			shading_type = context.space_data.shading.type
 			context.space_data.shading.type = wpt['shading_type']
 			if context.space_data.shading.type=='SOLID':
 				context.space_data.shading.light = wpt['light']
 				context.space_data.shading.color_type = wpt['color_type']
 				context.space_data.shading.single_color = wpt['single_color']
+			context.space_data.shading.type = shading_type
 
 		armature = get_armature_of_meshob(obj)
 		if not armature:
 			return {'FINISHED'}
 		# If an armature was un-hidden, hide it again.
-		armature.hide_viewport = wpt['armature_enabled']
-		armature.hide_set(wpt['armature_hide'])
-		armature.show_in_front = wpt['armature_in_front']
-		coll = bpy.data.collections.get(coll_name)
-		if coll:
-			bpy.data.collections.remove(coll)
+		armature.hide_viewport = wpt['arm_enabled']
+		armature.hide_set(wpt['arm_hide'])
+		armature.show_in_front = wpt['arm_in_front']
+
+		# Restore whether the armature is in local view or not.
+		if 'arm_local_view' in wpt and context.space_data.local_view:
+			armature.local_view_set(context.space_data, wpt['arm_local_view'])
+
+		# Remove armature from scene root collection if it was moved there.
+		if wpt['arm_coll_assigned']:
+			context.scene.collection.objects.unlink(armature)
+
 		return {'FINISHED'}
 
 	def execute(self, context):
@@ -165,15 +173,46 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 		else:
 			return self.leave_wp(context)
 
+class EASYWEIGHT_PT_Debug(bpy.types.Panel):
+	"""For debugging the Toggle Weight Paint operator."""
+	bl_space_type = 'VIEW_3D'
+	bl_region_type = 'UI'
+	bl_category = 'EasyWeight'
+	bl_idname = "EASYWEIGHT_PT_Debug"
+	bl_label = "Debug Toggle WP"
+
+	@classmethod
+	def poll(cls, context):
+		return context.preferences.addons[__package__].preferences.debug_mode
+
+	def draw(self, context):
+		layout = self.layout
+
+		if not 'wpt' in context.screen: 
+			return
+
+		wpt = context.screen['wpt'].to_dict()
+
+		for key in wpt.keys():
+			split = layout.split(factor=0.4)
+			value = wpt[key]
+			if type(value)==bpy.types.Object:
+				value = value.name
+			split.label(text=key)
+			split.label(text=str(value))
 
 def register():
 	from bpy.utils import register_class
 	register_class(EASYWEIGHT_OT_toggle_weight_paint)
+	register_class(EASYWEIGHT_PT_Debug)
+
 	bpy.types.VIEW3D_MT_paint_weight.append(EASYWEIGHT_OT_toggle_weight_paint.draw)
 	bpy.types.VIEW3D_MT_object.append(EASYWEIGHT_OT_toggle_weight_paint.draw)
 
 def unregister():
 	from bpy.utils import unregister_class
 	unregister_class(EASYWEIGHT_OT_toggle_weight_paint)
+	unregister_class(EASYWEIGHT_PT_Debug)
+
 	bpy.types.VIEW3D_MT_paint_weight.remove(EASYWEIGHT_OT_toggle_weight_paint.draw)
 	bpy.types.VIEW3D_MT_object.remove(EASYWEIGHT_OT_toggle_weight_paint.draw)