tweak_lattice.py

@@ -25,7 +25,7 @@ from .utils import clamp, get_lattice_vertex_index, simple_driver
 
 coll_name = 'Tweak Lattices'
 
-def ensure_tweak_lattice_collection(scene):
+def ensure_tweak_lattice_collection(scene: bpy.types.Scene) -> bpy.types.Collection:
 	coll = bpy.data.collections.get(coll_name)
 	if not coll:
 		coll = bpy.data.collections.new(coll_name)
@@ -33,7 +33,9 @@ def ensure_tweak_lattice_collection(scene):
 
 	return coll
 
-def ensure_falloff_vgroup(lattice_ob: bpy.types.Object, vg_name="Group", multiplier=1, power=1) -> bpy.types.VertexGroup:
+def ensure_falloff_vgroup(
+		lattice_ob: bpy.types.Object, 
+		vg_name="Group", multiplier=1, power=1) -> bpy.types.VertexGroup:
 	lattice = lattice_ob.data
 	res_x, res_y, res_z = lattice.points_u, lattice.points_v, lattice.points_w
 
@@ -62,6 +64,62 @@ def add_radius_constraint(obj, hook, target):
 		simple_driver(trans_con, prop, hook, '["Radius"]')
 	return trans_con
 
+def get_objects_of_lattice(hook: bpy.types.Object) -> List[bpy.types.Object]:
+	objs = []
+	ob_count = 0
+	ob_prop_name = "object_"+str(ob_count)
+	while ob_prop_name in hook:
+		objs.append(hook[ob_prop_name])
+	return objs
+
+def add_objects_to_lattice(
+		hook: bpy.types.Object, 
+		objects: List[bpy.types.Object]):
+	lattice_ob = hook['Lattice']
+	
+	# Check for existing
+	offset = 0
+	while "object_"+str(offset) in hook:
+		offset += 1
+
+	for i, o in enumerate(objects):
+		o.select_set(False)
+		if o.type!='MESH' or o in hook.values(): 
+			offset -= 1
+			continue
+		m = o.modifiers.new(name="Tweak Lattice", type='LATTICE')
+		m.object = lattice_ob
+		hook["object_"+str(i+offset)] = o
+		# Add driver to the modifier influence
+		simple_driver(m, 'strength', hook, '["Tweak Lattice"]')
+
+def remove_all_objects_from_lattice(hook: bpy.types.Object) -> List[bpy.types.Object]:
+	lattice_ob = hook['Lattice']
+	objs = []		
+	ob_count = 0
+	ob_prop_name = "object_"+str(ob_count)
+	while ob_prop_name in hook:
+		ob = hook[ob_prop_name]
+		for m in ob.modifiers:
+			if m.type!='LATTICE': continue
+			if m.object == lattice_ob:
+				m.driver_remove('strength')
+				ob.modifiers.remove(m)
+				break
+		ob_count += 1
+		objs.append(ob)
+		del hook[ob_prop_name]
+		ob_prop_name = "object_"+str(ob_count)
+	return objs
+
+def remove_objects_from_lattice(hook, objects):
+	new_objs = []
+	prev_objs = remove_all_objects_from_lattice(hook)
+	for o in prev_objs:
+		if o not in objects:
+			new_objs.append(o)
+	add_objects_to_lattice(hook, new_objs)
+
 class TWEAKLAT_OT_Create(bpy.types.Operator):
 	"""Create a lattice setup at the 3D cursor to deform selected objects"""
 	bl_idname = "lattice.create_tweak_lattice"
@@ -154,13 +212,7 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
 		hook_mod.vertex_group = vg.name
 
 		# Add Lattice modifier to the selected objects
-		for i, o in enumerate(context.selected_objects):
-			if o.type!='MESH': continue
-			m = o.modifiers.new(name="Tweak Lattice", type='LATTICE')
-			m.object = lattice_ob
-			hook["object_"+str(i)] = o
-			# Add driver to the modifier influence
-			simple_driver(m, 'strength', hook, '["Tweak Lattice"]')
+		add_objects_to_lattice(hook, context.selected_objects)
 		
 		# Set up Radius control.
 		add_radius_constraint(hook, hook, root)
@@ -243,18 +295,7 @@ class TWEAKLAT_OT_Delete(bpy.types.Operator):
 		root = hook['Root']
 
 		# Remove Lattice modifiers and their drivers.
-		ob_count = 0
-		ob_prop_name = "object_"+str(ob_count)
-		while ob_prop_name in hook:
-			ob = hook[ob_prop_name]
-			for m in ob.modifiers:
-				if m.type!='LATTICE': continue
-				if m.object == lattice:
-					m.driver_remove('strength')
-					ob.modifiers.remove(m)
-					break
-			ob_count += 1
-			ob_prop_name = "object_"+str(ob_count)
+		remove_all_objects_from_lattice(hook)
 		
 		# Remove hook Action if exists.
 		if hook.animation_data and hook.animation_data.action:
@@ -275,6 +316,57 @@ class TWEAKLAT_OT_Delete(bpy.types.Operator):
 
 		return {'FINISHED'}
 
+class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
+	"""Add selected objects to this tweak lattice"""
+	bl_idname = "lattice.add_selected_objects"
+	bl_label = "Add Selected Objects"
+	bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
+
+	@classmethod
+	def poll(cls, context):
+		ob = context.object
+		if ob.type!='EMPTY' or 'Tweak Lattice' not in ob: return False
+		values = ob.values()
+		for sel_o in context.selected_objects:
+			if sel_o==ob or sel_o.type!='MESH': continue
+			if sel_o not in values:
+				return True
+		return False
+
+	def execute(self, context):
+		hook = context.object
+
+		# Add Lattice modifier to the selected objects
+		add_objects_to_lattice(hook, context.selected_objects)
+
+		return {'FINISHED'}
+
+class TWEAKLAT_OT_Remove_Objects(bpy.types.Operator):
+	"""Remove selected objects from this tweak lattice"""
+	bl_idname = "lattice.remove_selected_objects"
+	bl_label = "Remove Selected Objects"
+	bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
+
+	@classmethod
+	def poll(cls, context):
+		ob = context.object
+		if ob.type!='EMPTY' or 'Tweak Lattice' not in ob: return False
+		values = ob.values()
+		for sel_o in context.selected_objects:
+			if sel_o==ob or sel_o.type!='MESH': continue
+			if sel_o in values:
+				return True
+		return False
+
+	def execute(self, context):
+		hook = context.object
+
+		# Add Lattice modifier to the selected objects
+		remove_objects_from_lattice(hook, context.selected_objects)
+
+		return {'FINISHED'}
+
+
 class TWEAKLAT_PT_Main(bpy.types.Panel):
 	bl_space_type = 'VIEW_3D'
 	bl_region_type = 'UI'
@@ -321,6 +413,22 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
 		root_row.prop(hook, '["Root"]', text="Root")
 		root_row.prop(hook['Root'], 'hide_viewport', text="", emboss=False)
 
+		layout.separator()
+		layout.label(text="Add Objects")
+		for o in context.selected_objects:
+			if o == hook or o.type!='MESH': continue
+			if o in hook.values(): continue
+			layout.label(text=o.name, icon='ADD')
+		layout.operator(TWEAKLAT_OT_Add_Objects.bl_idname, icon='ADD')
+
+		layout.separator()
+		layout.label(text="Remove Objects")
+		for o in context.selected_objects:
+			if o == hook or o.type!='MESH': continue
+			if o not in hook.values(): continue
+			layout.label(text=o.name, icon='REMOVE')
+		layout.operator(TWEAKLAT_OT_Remove_Objects.bl_idname, icon='REMOVE')
+
 		layout.separator()
 		layout.label(text="Affected Objects")
 
@@ -335,6 +443,8 @@ classes = [
 	TWEAKLAT_OT_Create
 	,TWEAKLAT_OT_Delete
 	,TWEAKLAT_OT_Falloff
+	,TWEAKLAT_OT_Add_Objects
+	,TWEAKLAT_OT_Remove_Objects
 	,TWEAKLAT_PT_Main
 ]
 
@@ -342,7 +452,7 @@ def register():
 	from bpy.utils import register_class
 	for c in classes:
 		register_class(c)
-
+	
 	bpy.types.Scene.tweak_lattice_radius = FloatProperty(name="Radius", default=0.1, min=0.0001, max=1000, soft_max=2)
 	bpy.types.Scene.tweak_lattice_parent_ob = PointerProperty(type=bpy.types.Object, name="Parent")	# Maybe it would be safer to make this a StringProperty but whatever.
 	bpy.types.Scene.tweak_lattice_parent_bone = StringProperty(name="Bone")