toggle_weight_paint.py

@@ -19,16 +19,139 @@ def get_armature_of_meshob(obj: bpy.types.Object):
 		if m.type=='ARMATURE':
 			return m.object
 
+def enter_wp(context) -> bool:
+	"""Enter weight paint mode, change the necessary settings, and save their
+	original states so they can be restored when leaving wp mode."""
+
+	obj = context.object
+	wm = context.window_manager
+
+	# Store old display_type setting in a Custom Property on the Object.
+	obj['wpt_display_type'] = obj.display_type
+	# Ensure display_type is SOLID; weights don't display in WIRE or BOUNDS.
+	if obj.display_type not in ['SOLID', 'TEXTURED']:
+		obj.display_type = 'SOLID'
+
+	# Store old shading settings in a Custom Property dictionary in the Scene.
+	if 'wpt' not in wm:
+		wm['wpt'] = {}
+
+	wpt = wm['wpt']
+	wpt_as_dict = wpt.to_dict()
+
+	# If we are entering WP mode for the first time or if the last time 
+	# the operator was exiting WP mode, then save current shading info.
+	if 'last_switch_in' not in wpt_as_dict or wpt_as_dict['last_switch_in']==False:
+		wpt['shading_type'] = context.space_data.shading.type
+		if context.space_data.shading.type=='SOLID':
+			# These properties only exist when the shading type is SOLID.
+			wpt['light'] = context.space_data.shading.light
+			wpt['color_type'] = context.space_data.shading.color_type
+			wpt['single_color'] = context.space_data.shading.single_color
+		wpt['active_object'] = obj
+
+	# This flag indicates that the last time this operator ran, we were 
+	# switching INTO wp mode.
+	wpt['last_switch_in'] = True
+	wpt['mode'] = obj.mode
+
+	# Set shading
+	if context.space_data.shading.type=='SOLID':
+		context.space_data.shading.light = 'FLAT'
+		context.space_data.shading.color_type = 'SINGLE'
+		context.space_data.shading.single_color = (1,1,1)
+
+	# Enter WP mode.
+	bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
+
+	### ENSURING ARMATURE VISIBILITY
+	armature = get_armature_of_meshob(obj)
+	if not armature:
+		return
+	# Save all object visibility related info so it can be restored later.
+	wpt['arm_enabled'] = armature.hide_viewport
+	wpt['arm_hide'] = armature.hide_get()
+	wpt['arm_in_front'] = armature.show_in_front
+	wpt['arm_coll_assigned'] = False
+	armature.hide_viewport = False
+	armature.hide_set(False)
+	armature.show_in_front = True
+	if context.space_data.local_view:
+		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, 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)
+		wpt['arm_coll_assigned'] = True
+
+	if armature.visible_get():
+		context.view_layer.objects.active = armature
+		bpy.ops.object.mode_set(mode='POSE')
+	
+	context.view_layer.objects.active = obj
+	return armature.visible_get()
+
+def leave_wp(context):
+	"""Leave weight paint mode, then find, restore, and delete the data
+	that was stored about shading settings in enter_wp()."""
+
+	obj = context.object
+	wm = context.window_manager
+
+	# Restore object display type
+	if 'wpt_display_type' in obj:
+		obj.display_type = obj['wpt_display_type']
+		del obj['wpt_display_type']
+
+	if 'wpt' not in wm or 'mode' not in wm['wpt'].to_dict():
+		# There is no saved data to restore from, nothing else to do.
+		bpy.ops.object.mode_set(mode='OBJECT')
+		return {'FINISHED'}
+
+	wpt = wm['wpt']
+	wpt_as_dict = wpt.to_dict()
+
+	# Restore mode.
+	bpy.ops.object.mode_set(mode=wpt_as_dict['mode'])
+
+	# Reset the stored data
+	wm['wpt'] = {}
+	# Flag to save that the last time the operator ran we were EXITING wp mode.
+	wm['wpt']['last_switch_in'] = False
+
+	# Restore shading options.
+	if 'light' in wpt_as_dict:	# 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_as_dict['shading_type']
+		if context.space_data.shading.type=='SOLID':
+			context.space_data.shading.light = wpt_as_dict['light']
+			context.space_data.shading.color_type = wpt_as_dict['color_type']
+			context.space_data.shading.single_color = wpt_as_dict['single_color']
+		context.space_data.shading.type = shading_type
+
+	armature = get_armature_of_meshob(obj)
+	if not armature:
+		return
+	# If an armature was un-hidden, hide it again.
+	armature.hide_viewport = wpt_as_dict['arm_enabled']
+	armature.hide_set(wpt_as_dict['arm_hide'])
+	armature.show_in_front = wpt_as_dict['arm_in_front']
+
+	# Restore whether the armature is in local view or not.
+	if 'arm_local_view' in wpt_as_dict and context.space_data.local_view:
+		armature.local_view_set(context.space_data, wpt_as_dict['arm_local_view'])
+
+	# Remove armature from scene root collection if it was moved there.
+	if wpt_as_dict['arm_coll_assigned']:
+		context.scene.collection.objects.unlink(armature)
+
+	return
 class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 	"""Toggle weight paint mode properly with a single operator. """
 	bl_idname = "object.weight_paint_toggle"
 	bl_label = "Toggle Weight Paint Mode"
 	bl_options = {'REGISTER', 'UNDO'}
-
-	# local_view: bpy.props.BoolProperty(
-	# 	name		 = "Local View"
-	# 	,description = "Enter Local view with the mesh and armature"
-	# )
 	
 	@classmethod
 	def poll(cls, context):
@@ -38,144 +161,18 @@ class EASYWEIGHT_OT_toggle_weight_paint(bpy.types.Operator):
 	def draw(self, context):
 		self.layout.operator(EASYWEIGHT_OT_toggle_weight_paint.bl_idname)
 
-	def enter_wp(self, context):
-		"""Enter weight paint mode, change the necessary settings, and save their
-		original states so they can be restored when leaving wp mode."""
-
-		obj = context.object
-		wm = context.window_manager
-
-		# Store old display_type setting in a Custom Property on the Object.
-		obj['wpt_display_type'] = obj.display_type
-		# Ensure display_type is SOLID; weights don't display in WIRE or BOUNDS.
-		if obj.display_type not in ['SOLID', 'TEXTURED']:
-			obj.display_type = 'SOLID'
-
-		# Store old shading settings in a Custom Property dictionary in the Scene.
-		if 'wpt' not in wm:
-			wm['wpt'] = {}
-
-		wpt = wm['wpt']
-		wpt_as_dict = wpt.to_dict()
-
-		# If we are entering WP mode for the first time or if the last time 
-		# the operator was exiting WP mode, then save current shading info.
-		if 'last_switch_in' not in wpt_as_dict or wpt_as_dict['last_switch_in']==False:
-			wpt['shading_type'] = context.space_data.shading.type
-			if context.space_data.shading.type=='SOLID':
-				# These properties only exist when the shading type is SOLID.
-				wpt['light'] = context.space_data.shading.light
-				wpt['color_type'] = context.space_data.shading.color_type
-				wpt['single_color'] = context.space_data.shading.single_color
-			wpt['active_object'] = obj
-
-		# This flag indicates that the last time this operator ran, we were 
-		# switching INTO wp mode.
-		wpt['last_switch_in'] = True
-		wpt['mode'] = obj.mode
-
-		# Set shading
-		if context.space_data.shading.type=='SOLID':
-			context.space_data.shading.light = 'FLAT'
-			context.space_data.shading.color_type = 'SINGLE'
-			context.space_data.shading.single_color = (1,1,1)
-
-		# Enter WP mode.
-		bpy.ops.object.mode_set(mode='WEIGHT_PAINT')
-
-		### ENSURING ARMATURE VISIBILITY
-		armature = get_armature_of_meshob(obj)
-		if not armature:
-			return {'FINISHED'}
-		# Save all object visibility related info so it can be restored later.
-		wpt['arm_enabled'] = armature.hide_viewport
-		wpt['arm_hide'] = armature.hide_get()
-		wpt['arm_in_front'] = armature.show_in_front
-		wpt['arm_coll_assigned'] = False
-		armature.hide_viewport = False
-		armature.hide_set(False)
-		armature.show_in_front = True
-		if context.space_data.local_view:
-			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, 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)
-			wpt['arm_coll_assigned'] = True
-
-		if not armature.visible_get():
-			self.report({'WARNING'}, "Failed to make the armature visible. Curse you, Blender!!!")
-		else:
-			context.view_layer.objects.active = armature
-			bpy.ops.object.mode_set(mode='POSE')
-		
-		context.view_layer.objects.active = obj
-		return {'FINISHED'}
-
-	def leave_wp(self, context):
-		"""Leave weight paint mode, then find, restore, and delete the data
-		that was stored about shading settings in enter_wp()."""
-
-		obj = context.object
-		wm = context.window_manager
-
-		# Restore object display type
-		if 'wpt_display_type' in obj:
-			obj.display_type = obj['wpt_display_type']
-			del obj['wpt_display_type']
-
-		if 'wpt' not in wm or 'mode' not in wm['wpt'].to_dict():
-			# There is no saved data to restore from, nothing else to do.
-			bpy.ops.object.mode_set(mode='OBJECT')
-			return {'FINISHED'}
-
-		wpt = wm['wpt']
-		wpt_as_dict = wpt.to_dict()
-
-		# Restore mode.
-		bpy.ops.object.mode_set(mode=wpt_as_dict['mode'])
-
-		# Reset the stored data
-		wm['wpt'] = {}
-		# Flag to save that the last time the operator ran we were EXITING wp mode.
-		wm['wpt']['last_switch_in'] = False
-
-		# Restore shading options.
-		if 'light' in wpt_as_dict:	# 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_as_dict['shading_type']
-			if context.space_data.shading.type=='SOLID':
-				context.space_data.shading.light = wpt_as_dict['light']
-				context.space_data.shading.color_type = wpt_as_dict['color_type']
-				context.space_data.shading.single_color = wpt_as_dict['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_as_dict['arm_enabled']
-		armature.hide_set(wpt_as_dict['arm_hide'])
-		armature.show_in_front = wpt_as_dict['arm_in_front']
-
-		# Restore whether the armature is in local view or not.
-		if 'arm_local_view' in wpt_as_dict and context.space_data.local_view:
-			armature.local_view_set(context.space_data, wpt_as_dict['arm_local_view'])
-
-		# Remove armature from scene root collection if it was moved there.
-		if wpt_as_dict['arm_coll_assigned']:
-			context.scene.collection.objects.unlink(armature)
-
-		return {'FINISHED'}
-
 	def execute(self, context):
 		obj = context.object
 
 		if obj.mode != 'WEIGHT_PAINT':
-			return self.enter_wp(context)
+			armature_visible = enter_wp(context)
+			if armature_visible == False:
+				# This should never happen, but it also doesn't break anything.
+				self.report({'WARNING'}, "Could not make Armature visible.")
+			return {'FINISHED'}
 		else:
-			return self.leave_wp(context)
+			leave_wp(context)
+			return {'FINISHED'}
 
 class EASYWEIGHT_PT_Debug(bpy.types.Panel):
 	"""For debugging the Toggle Weight Paint operator."""