tweak_lattice.py

@@ -11,115 +11,8 @@ from rna_prop_ui import rna_idprop_ui_create
 
 from .utils import clamp, get_lattice_vertex_index, simple_driver, bounding_box_center_of_objects
 
-coll_name = 'Tweak Lattices'
+TWEAKLAT_COLL_NAME = 'Tweak Lattices'
 
-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)
-		scene.collection.children.link(coll)
-
-	return coll
-
-def ensure_falloff_vgroup(
-		lattice_ob: bpy.types.Object, 
-		vg_name="Group", multiplier=1, expression="x") -> bpy.types.VertexGroup:
-	lattice = lattice_ob.data
-	res_x, res_y, res_z = lattice.points_u, lattice.points_v, lattice.points_w
-
-	vg = lattice_ob.vertex_groups.get(vg_name)
-
-	center = Vector((res_x-1, res_y-1, res_z-1))/2
-	max_res = max(res_x, res_y, res_z)
-
-	if not vg:
-		vg = lattice_ob.vertex_groups.new(name=vg_name)
-	for x in range(res_x-4):
-		for y in range(res_y-4):
-			for z in range(res_z-4):
-				index = get_lattice_vertex_index(lattice, (x+2, y+2, z+2))
-
-				coord = Vector((x+2, y+2, z+2))
-				distance_from_center = (coord-center).length
-				distance_factor = 1 - (distance_from_center / max_res * 2)
-				try:
-					influence = eval(expression.replace("x", "distance_factor"))
-				except:
-					return None
-
-				vg.add([index], influence * multiplier, 'REPLACE')
-	return vg
-
-def add_radius_constraint(obj, hook, target):
-	trans_con = obj.constraints.new(type='TRANSFORM')
-	trans_con.target = target
-	trans_con.map_to = 'SCALE'
-	trans_con.mix_mode_scale = 'MULTIPLY'
-	for prop in ['to_min_x_scale', 'to_min_y_scale', 'to_min_z_scale']:
-		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 get_lattice_modifier_of_object(obj, lattice) -> bpy.types.Modifier:
-	"""Find the lattice modifier on the object that uses this lattice"""
-	for m in obj.modifiers:
-		if m.type == 'LATTICE' and m.object == lattice:
-			return m
-
-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 to deform selected objects"""
@@ -135,9 +28,9 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
     )
 
     location: EnumProperty(name="Location", items=[
-		('CURSOR', "3D Cursor", "Create at the location and orientation of the 3D cursor.")
+        ('CURSOR', "3D Cursor", "Create at the location and orientation of the 3D cursor."),
-		,('CENTER', "Center", "Create at the bounding box center of all selected objects.")
+        ('CENTER', "Center", "Create at the bounding box center of all selected objects."),
-		,('PARENT', "Parent", "Create at the location of the parent object or bone.")
+        ('PARENT', "Parent", "Create at the location of the parent object or bone.")
     ])
     radius: FloatProperty(
         name="Radius",
@@ -147,7 +40,8 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         max=1000,
         soft_max=2
     )
-	parent_bone: StringProperty(name="Bone", description="Bone to use as parent")
+    parent_bone: StringProperty(
+        name="Bone", description="Bone to use as parent")
 
     @classmethod
     def poll(cls, context):
@@ -184,8 +78,8 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
 
         scene = context.scene
         if scene.tweak_lattice_parent_ob and scene.tweak_lattice_parent_ob.type == 'ARMATURE':
-			col.prop_search(self, 'parent_bone', scene.tweak_lattice_parent_ob.data, 'bones')
+            col.prop_search(self, 'parent_bone',
-
+                            scene.tweak_lattice_parent_ob.data, 'bones')
 
     def execute(self, context):
         scene = context.scene
@@ -240,13 +134,15 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         root.empty_display_size = 0.5
         if self.location == 'CENTER':
             meshes = [o for o in context.selected_objects if o.type == 'MESH']
-			root.matrix_world.translation = bounding_box_center_of_objects(meshes)
+            root.matrix_world.translation = bounding_box_center_of_objects(
+                meshes)
         elif self.location == 'CURSOR':
             root.matrix_world = context.scene.cursor.matrix
         elif self.location == 'PARENT':
             matrix_of_parent = scene.tweak_lattice_parent_ob.matrix_world
             if self.parent_bone:
-				matrix_of_parent = scene.tweak_lattice_parent_ob.matrix_world @ scene.tweak_lattice_parent_ob.pose.bones[self.parent_bone].matrix
+                matrix_of_parent = scene.tweak_lattice_parent_ob.matrix_world @ scene.tweak_lattice_parent_ob.pose.bones[
+                    self.parent_bone].matrix
             root.matrix_world = matrix_of_parent.copy()
         coll.objects.link(root)
         root.hide_viewport = True
@@ -280,7 +176,8 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         add_radius_constraint(hook, hook, root)
         add_radius_constraint(lattice_ob, hook, root)
 
-		root_drv = simple_driver(root, 'empty_display_size', hook, '["Radius"]')
+        root_drv = simple_driver(
+            root, 'empty_display_size', hook, '["Radius"]')
         root_drv.expression = 'var/2'
 
         # Deselect everything, select the hook and make it active
@@ -291,6 +188,7 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         scene.tweak_lattice_parent_ob = None
         return {'FINISHED'}
 
+
 class TWEAKLAT_OT_Falloff(bpy.types.Operator):
     """Adjust falloff of the hook vertex group of a Tweak Lattice"""
     bl_idname = "lattice.tweak_lattice_adjust_falloww"
@@ -298,10 +196,12 @@ class TWEAKLAT_OT_Falloff(bpy.types.Operator):
     bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
 
     def update_falloff(self, context):
-		if self.doing_invoke: return
+        if self.doing_invoke:
+            return
         hook = context.object
         lattice_ob = hook['Lattice']
-		ret = ensure_falloff_vgroup(lattice_ob, 'Hook', multiplier=self.multiplier, expression=self.expression)
+        ret = ensure_falloff_vgroup(
+            lattice_ob, 'Hook', multiplier=self.multiplier, expression=self.expression)
         self.is_expression_valid = ret != None
         if ret:
             hook['Expression'] = self.expression
@@ -368,6 +268,7 @@ class TWEAKLAT_OT_Falloff(bpy.types.Operator):
     def execute(self, context):
         return {'FINISHED'}
 
+
 class TWEAKLAT_OT_Delete(bpy.types.Operator):
     """Delete a tweak lattice setup with all its helper objects, drivers, etc"""
     bl_idname = "lattice.delete_tweak_lattice"
@@ -399,13 +300,13 @@ class TWEAKLAT_OT_Delete(bpy.types.Operator):
         bpy.data.objects.remove(root)
 
         # Remove the collection if it's empty.
-		coll = bpy.data.collections.get(coll_name)
+        coll = bpy.data.collections.get(TWEAKLAT_COLL_NAME)
         if coll and len(coll.all_objects) == 0:
             bpy.data.collections.remove(coll)
 
-
         return {'FINISHED'}
 
+
 class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
     """Add selected objects to this tweak lattice"""
     bl_idname = "lattice.add_selected_objects"
@@ -415,10 +316,12 @@ class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
     @classmethod
     def poll(cls, context):
         ob = context.object
-		if ob.type!='EMPTY' or 'Tweak Lattice' not in ob: return False
+        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 == ob or sel_o.type != 'MESH':
+                continue
             if sel_o not in values:
                 return True
         return False
@@ -431,6 +334,7 @@ class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
 
         return {'FINISHED'}
 
+
 class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
     """Remove selected objects from this tweak lattice"""
     bl_idname = "lattice.remove_selected_objects"
@@ -440,10 +344,12 @@ class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
     @classmethod
     def poll(cls, context):
         ob = context.object
-		if ob.type!='EMPTY' or 'Tweak Lattice' not in ob: return False
+        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 == ob or sel_o.type != 'MESH':
+                continue
             if sel_o in values:
                 return True
         return False
@@ -456,13 +362,15 @@ class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
 
         return {'FINISHED'}
 
+
 class TWEAKLAT_OT_Remove_Object(bpy.types.Operator):
     """Remove this object from the tweak lattice"""
     bl_idname = "lattice.remove_object"
     bl_label = "Remove Object"
     bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
 
-	ob_pointer_prop_name: StringProperty(description="Name of the custom property that references the object that should be removed")
+    ob_pointer_prop_name: StringProperty(
+        description="Name of the custom property that references the object that should be removed")
 
     def execute(self, context):
         hook = context.object
@@ -473,6 +381,7 @@ class TWEAKLAT_OT_Remove_Object(bpy.types.Operator):
 
         return {'FINISHED'}
 
+
 class TWEAKLAT_PT_Main(bpy.types.Panel):
     bl_space_type = 'VIEW_3D'
     bl_region_type = 'UI'
@@ -491,7 +400,8 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
         hook = context.object
         layout = layout.column()
         if hook.type != 'EMPTY' or 'Tweak Lattice' not in hook:
-			layout.operator(TWEAKLAT_OT_Create.bl_idname, icon='OUTLINER_OB_LATTICE')
+            layout.operator(TWEAKLAT_OT_Create.bl_idname,
+                            icon='OUTLINER_OB_LATTICE')
             return
 
         layout.prop(hook, '["Tweak Lattice"]', slider=True, text="Influence")
@@ -499,13 +409,15 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
         layout.operator(TWEAKLAT_OT_Falloff.bl_idname, text="Adjust Falloff")
 
         layout.separator()
-		layout.operator(TWEAKLAT_OT_Delete.bl_idname, text='Delete Tweak Lattice', icon='TRASH')
+        layout.operator(TWEAKLAT_OT_Delete.bl_idname,
+                        text='Delete Tweak Lattice', icon='TRASH')
 
         layout.separator()
         layout.label(text="Helper Objects")
         lattice_row = layout.row()
         lattice_row.prop(hook, '["Lattice"]', text="Lattice")
-		lattice_row.prop(hook['Lattice'], 'hide_viewport', text="", emboss=False)
+        lattice_row.prop(hook['Lattice'], 'hide_viewport',
+                         text="", emboss=False)
 
         root_row = layout.row()
         root_row.prop(hook, '["Root"]', text="Root")
@@ -525,56 +437,180 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
 
         num_to_add = 0
         for o in context.selected_objects:
-			if o == hook or o.type!='MESH': continue
+            if o == hook or o.type != 'MESH':
-			if o in hook.values(): continue
+                continue
+            if o in hook.values():
+                continue
             num_to_add += 1
             if num_to_add == 1:
                 text = f"Add {o.name}"
         if num_to_add:
             if num_to_add > 1:
                 text = f"Add {num_to_add} Objects"
-			layout.operator(TWEAKLAT_OT_Add_Objects.bl_idname, icon='ADD', text=text)
+            layout.operator(TWEAKLAT_OT_Add_Objects.bl_idname,
+                            icon='ADD', text=text)
 
         layout.separator()
         num_to_remove = False
         for o in context.selected_objects:
-			if o == hook or o.type!='MESH': continue
+            if o == hook or o.type != 'MESH':
-			if o not in hook.values(): continue
+                continue
+            if o not in hook.values():
+                continue
             num_to_remove += 1
             if num_to_remove == 1:
                 text = f"Remove {o.name}"
         if num_to_remove:
             if num_to_remove > 1:
                 text = f"Remove {num_to_remove} Objects"
-			layout.operator(TWEAKLAT_OT_Remove_Selected_Objects.bl_idname, icon='REMOVE', text=text)
+            layout.operator(
+                TWEAKLAT_OT_Remove_Selected_Objects.bl_idname, icon='REMOVE', text=text)
 
-		objects_and_keys = [(hook[key], key) for key in hook.keys() if "object_" in key]
+        objects_and_keys = [(hook[key], key)
+                            for key in hook.keys() if "object_" in key]
         objects_and_keys.sort(key=lambda o_and_k: o_and_k[1])
         for ob, key in objects_and_keys:
             row = layout.row(align=True)
             row.prop(hook, f'["{key}"]', text="")
             mod = get_lattice_modifier_of_object(ob, hook['Lattice'])
-			row.prop_search(mod, 'vertex_group', ob, 'vertex_groups', text="", icon='GROUP_VERTEX')
+            row.prop_search(mod, 'vertex_group', ob,
-			op = row.operator(TWEAKLAT_OT_Remove_Object.bl_idname, text="", icon='X')
+                            'vertex_groups', text="", icon='GROUP_VERTEX')
+            op = row.operator(
+                TWEAKLAT_OT_Remove_Object.bl_idname, text="", icon='X')
             op.ob_pointer_prop_name = key
 
 
+def ensure_tweak_lattice_collection(scene: bpy.types.Scene) -> bpy.types.Collection:
+    coll = bpy.data.collections.get(TWEAKLAT_COLL_NAME)
+    if not coll:
+        coll = bpy.data.collections.new(TWEAKLAT_COLL_NAME)
+        scene.collection.children.link(coll)
+
+    return coll
+
+
+def ensure_falloff_vgroup(
+        lattice_ob: bpy.types.Object,
+        vg_name="Group", multiplier=1, expression="x") -> bpy.types.VertexGroup:
+    lattice = lattice_ob.data
+    res_x, res_y, res_z = lattice.points_u, lattice.points_v, lattice.points_w
+
+    vg = lattice_ob.vertex_groups.get(vg_name)
+
+    center = Vector((res_x-1, res_y-1, res_z-1))/2
+    max_res = max(res_x, res_y, res_z)
+
+    if not vg:
+        vg = lattice_ob.vertex_groups.new(name=vg_name)
+    for x in range(res_x-4):
+        for y in range(res_y-4):
+            for z in range(res_z-4):
+                index = get_lattice_vertex_index(lattice, (x+2, y+2, z+2))
+
+                coord = Vector((x+2, y+2, z+2))
+                distance_from_center = (coord-center).length
+                distance_factor = 1 - (distance_from_center / max_res * 2)
+                try:
+                    influence = eval(expression.replace(
+                        "x", "distance_factor"))
+                except:
+                    return None
+
+                vg.add([index], influence * multiplier, 'REPLACE')
+    return vg
+
+
+def add_radius_constraint(obj, hook, target):
+    trans_con = obj.constraints.new(type='TRANSFORM')
+    trans_con.target = target
+    trans_con.map_to = 'SCALE'
+    trans_con.mix_mode_scale = 'MULTIPLY'
+    for prop in ['to_min_x_scale', 'to_min_y_scale', 'to_min_z_scale']:
+        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 get_lattice_modifier_of_object(obj, lattice) -> bpy.types.Modifier:
+    """Find the lattice modifier on the object that uses this lattice"""
+    for m in obj.modifiers:
+        if m.type == 'LATTICE' and m.object == lattice:
+            return m
+
+
+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)
+
+
 classes = [
-	TWEAKLAT_OT_Create
+    TWEAKLAT_OT_Create, TWEAKLAT_OT_Delete, TWEAKLAT_OT_Falloff, TWEAKLAT_OT_Add_Objects, TWEAKLAT_OT_Remove_Selected_Objects, TWEAKLAT_OT_Remove_Object, TWEAKLAT_PT_Main
-	,TWEAKLAT_OT_Delete
-	,TWEAKLAT_OT_Falloff
-	,TWEAKLAT_OT_Add_Objects
-	,TWEAKLAT_OT_Remove_Selected_Objects
-	,TWEAKLAT_OT_Remove_Object
-	,TWEAKLAT_PT_Main
 ]
 
+
 def register():
     from bpy.utils import register_class
     for c in classes:
         register_class(c)
 
-	bpy.types.Scene.tweak_lattice_parent_ob = PointerProperty(type=bpy.types.Object, name="Parent")
+    bpy.types.Scene.tweak_lattice_parent_ob = PointerProperty(
+        type=bpy.types.Object, name="Parent")
+
 
 def unregister():
     from bpy.utils import unregister_class