tweak_lattice.py

@@ -5,7 +5,9 @@
 
 import bpy
 from bpy.props import FloatProperty, IntVectorProperty, FloatVectorProperty, BoolProperty, PointerProperty, StringProperty, EnumProperty
-from typing import List
+from bpy.types import Operator, Object, VertexGroup, Scene, Collection, Modifier, Panel
+from typing import List, Tuple
+
 from mathutils import Vector
 from rna_prop_ui import rna_idprop_ui_create
 
@@ -14,7 +16,7 @@ from .utils import clamp, get_lattice_vertex_index, simple_driver, bounding_box_
 TWEAKLAT_COLL_NAME = 'Tweak Lattices'
 
 
-class TWEAKLAT_OT_Create(bpy.types.Operator):
+class TWEAKLAT_OT_Create(Operator):
     """Create a lattice setup to deform selected objects"""
     bl_idname = "lattice.create_tweak_lattice"
     bl_label = "Create Tweak Lattice"
@@ -112,6 +114,7 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
 
         # Create some custom properties
         hook['Lattice'] = lattice_ob
+        lattice_ob['Hook'] = hook
         hook['Multiplier'] = 1.0
         hook['Expression'] = 'x'
 
@@ -130,6 +133,7 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         # This will allow pressing Ctrl+G/R/S on the hook to reset its transforms.
         root_name = "Root_" + hook.name
         root = bpy.data.objects.new(root_name, None)
+        root['Hook'] = hook
         root.empty_display_type = 'CUBE'
         root.empty_display_size = 0.5
         if self.location == 'CENTER':
@@ -189,7 +193,7 @@ class TWEAKLAT_OT_Create(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_OT_Falloff(bpy.types.Operator):
+class TWEAKLAT_OT_Falloff(Operator):
     """Adjust falloff of the hook vertex group of a Tweak Lattice"""
     bl_idname = "lattice.tweak_lattice_adjust_falloww"
     bl_label = "Adjust Falloff"
@@ -198,10 +202,9 @@ class TWEAKLAT_OT_Falloff(bpy.types.Operator):
     def update_falloff(self, context):
         if self.doing_invoke:
             return
-        hook = context.object
+        hook, lattice, _root = get_tweak_setup(context.object)
-        lattice_ob = hook['Lattice']
         ret = ensure_falloff_vgroup(
-            lattice_ob, 'Hook', multiplier=self.multiplier, expression=self.expression)
+            lattice, 'Hook', multiplier=self.multiplier, expression=self.expression)
         self.is_expression_valid = ret != None
         if ret:
             hook['Expression'] = self.expression
@@ -235,11 +238,11 @@ class TWEAKLAT_OT_Falloff(bpy.types.Operator):
 
     @classmethod
     def poll(cls, context):
-        ob = context.object
+        hook, lattice, root = get_tweak_setup(context.object)
-        return ob.type == 'EMPTY' and 'Tweak Lattice' in ob
+        return hook and lattice and root
 
     def invoke(self, context, event):
-        hook = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
         self.multiplier = hook['Multiplier']
         self.hook_start_scale = hook.scale.copy()
         lattice_ob = hook['Lattice']
@@ -269,7 +272,7 @@ class TWEAKLAT_OT_Falloff(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_OT_Delete(bpy.types.Operator):
+class TWEAKLAT_OT_Delete(Operator):
     """Delete a tweak lattice setup with all its helper objects, drivers, etc"""
     bl_idname = "lattice.delete_tweak_lattice"
     bl_label = "Delete Tweak Lattice"
@@ -277,13 +280,11 @@ class TWEAKLAT_OT_Delete(bpy.types.Operator):
 
     @classmethod
     def poll(cls, context):
-        ob = context.object
+        hook, lattice, root = get_tweak_setup(context.object)
-        return ob and ob.type == 'EMPTY' and 'Tweak Lattice' in ob
+        return hook and lattice and root
 
     def execute(self, context):
-        hook = context.object
+        hook, lattice, root = get_tweak_setup(context.object)
-        lattice = hook['Lattice']
-        root = hook['Root']
 
         # Remove Lattice modifiers and their drivers.
         remove_all_objects_from_lattice(hook)
@@ -307,7 +308,7 @@ class TWEAKLAT_OT_Delete(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
+class TWEAKLAT_OT_Add_Objects(Operator):
     """Add selected objects to this tweak lattice"""
     bl_idname = "lattice.add_selected_objects"
     bl_label = "Add Selected Objects"
@@ -315,19 +316,20 @@ class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
 
     @classmethod
     def poll(cls, context):
-        ob = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
-        if ob.type != 'EMPTY' or 'Tweak Lattice' not in ob:
+        if not hook:
             return False
-        values = ob.values()
+
+        values = hook.values()
         for sel_o in context.selected_objects:
-            if sel_o == ob or sel_o.type != 'MESH':
+            if sel_o == hook or sel_o.type != 'MESH':
                 continue
             if sel_o not in values:
                 return True
         return False
 
     def execute(self, context):
-        hook = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
 
         # Add Lattice modifier to the selected objects
         add_objects_to_lattice(hook, context.selected_objects)
@@ -335,7 +337,7 @@ class TWEAKLAT_OT_Add_Objects(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
+class TWEAKLAT_OT_Remove_Selected_Objects(Operator):
     """Remove selected objects from this tweak lattice"""
     bl_idname = "lattice.remove_selected_objects"
     bl_label = "Remove Selected Objects"
@@ -343,19 +345,20 @@ class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
 
     @classmethod
     def poll(cls, context):
-        ob = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
-        if ob.type != 'EMPTY' or 'Tweak Lattice' not in ob:
+        if not hook:
             return False
-        values = ob.values()
+
+        values = hook.values()
         for sel_o in context.selected_objects:
-            if sel_o == ob or sel_o.type != 'MESH':
+            if sel_o == hook or sel_o.type != 'MESH':
                 continue
             if sel_o in values:
                 return True
         return False
 
     def execute(self, context):
-        hook = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
 
         # Add Lattice modifier to the selected objects
         remove_objects_from_lattice(hook, context.selected_objects)
@@ -363,7 +366,7 @@ class TWEAKLAT_OT_Remove_Selected_Objects(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_OT_Remove_Object(bpy.types.Operator):
+class TWEAKLAT_OT_Remove_Object(Operator):
     """Remove this object from the tweak lattice"""
     bl_idname = "lattice.remove_object"
     bl_label = "Remove Object"
@@ -373,7 +376,7 @@ class TWEAKLAT_OT_Remove_Object(bpy.types.Operator):
         description="Name of the custom property that references the object that should be removed")
 
     def execute(self, context):
-        hook = context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
         target = hook[self.ob_pointer_prop_name]
 
         # Add Lattice modifier to the selected objects
@@ -382,7 +385,17 @@ class TWEAKLAT_OT_Remove_Object(bpy.types.Operator):
         return {'FINISHED'}
 
 
-class TWEAKLAT_PT_Main(bpy.types.Panel):
+class TWEAKLAT_OT_Duplicate(Operator):
+    """Duplicate this Tweak Lattice set-up"""
+    bl_idname = "lattice.duplicate_tweak_setup"
+    bl_label = "Duplicate Tweak Lattice"
+    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
+
+    def execute(self, context):
+        pass
+
+
+class TWEAKLAT_PT_Main(Panel):
     bl_space_type = 'VIEW_3D'
     bl_region_type = 'UI'
     bl_category = 'Lattice Magic'
@@ -390,16 +403,19 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
 
     @classmethod
     def poll(cls, context):
-        return context.object
+        hook, _lattice, _root = get_tweak_setup(context.object)
+
+        return context.object.type == 'MESH' or hook
 
     def draw(self, context):
         layout = self.layout
         layout.use_property_split = True
         layout.use_property_decorate = False
 
-        hook = context.object
+        hook, lattice, root = get_tweak_setup(context.object)
+
         layout = layout.column()
-        if hook.type != 'EMPTY' or 'Tweak Lattice' not in hook:
+        if not hook:
             layout.operator(TWEAKLAT_OT_Create.bl_idname,
                             icon='OUTLINER_OB_LATTICE')
             return
@@ -416,16 +432,15 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
         layout.label(text="Helper Objects")
         lattice_row = layout.row()
         lattice_row.prop(hook, '["Lattice"]', text="Lattice")
-        lattice_row.prop(hook['Lattice'], 'hide_viewport',
+        lattice_row.prop(lattice, 'hide_viewport',
                          text="", emboss=False)
 
         root_row = layout.row()
         root_row.prop(hook, '["Root"]', text="Root")
-        root_row.prop(hook['Root'], 'hide_viewport', text="", emboss=False)
+        root_row.prop(root, 'hide_viewport', text="", emboss=False)
 
         layout.separator()
         layout.label(text="Parenting")
-        root = hook['Root']
         col = layout.column()
         col.enabled = False
         col.prop(root, 'parent')
@@ -472,7 +487,7 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
         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'])
+            mod = get_lattice_modifier_of_object(ob, lattice)
             row.prop_search(mod, 'vertex_group', ob,
                             'vertex_groups', text="", icon='GROUP_VERTEX')
             op = row.operator(
@@ -480,7 +495,20 @@ class TWEAKLAT_PT_Main(bpy.types.Panel):
             op.ob_pointer_prop_name = key
 
 
-def ensure_tweak_lattice_collection(scene: bpy.types.Scene) -> bpy.types.Collection:
+def get_tweak_setup(obj: Object) -> Tuple[Object, Object, Object]:
+    """Based on either the hook, lattice or root, return all three."""
+    if obj.type == 'EMPTY':
+        if 'Root' and 'Lattice' in obj:
+            return obj, obj['Lattice'], obj['Root']
+        elif 'Hook' in obj:
+            return obj['Hook'], obj['Hook']['Lattice'], obj
+    elif obj.type == 'LATTICE' and 'Hook' in obj:
+        return obj['Hook'], obj, obj['Hook']['Root']
+
+    return [None, None, None]
+
+
+def ensure_tweak_lattice_collection(scene: Scene) -> Collection:
     coll = bpy.data.collections.get(TWEAKLAT_COLL_NAME)
     if not coll:
         coll = bpy.data.collections.new(TWEAKLAT_COLL_NAME)
@@ -490,8 +518,8 @@ def ensure_tweak_lattice_collection(scene: bpy.types.Scene) -> bpy.types.Collect
 
 
 def ensure_falloff_vgroup(
-        lattice_ob: bpy.types.Object,
+        lattice_ob: Object,
-        vg_name="Group", multiplier=1, expression="x") -> bpy.types.VertexGroup:
+        vg_name="Group", multiplier=1, expression="x") -> VertexGroup:
     lattice = lattice_ob.data
     res_x, res_y, res_z = lattice.points_u, lattice.points_v, lattice.points_w
 
@@ -530,7 +558,7 @@ def add_radius_constraint(obj, hook, target):
     return trans_con
 
 
-def get_objects_of_lattice(hook: bpy.types.Object) -> List[bpy.types.Object]:
+def get_objects_of_lattice(hook: Object) -> List[Object]:
     objs = []
     ob_count = 0
     ob_prop_name = "object_"+str(ob_count)
@@ -539,7 +567,7 @@ def get_objects_of_lattice(hook: bpy.types.Object) -> List[bpy.types.Object]:
     return objs
 
 
-def get_lattice_modifier_of_object(obj, lattice) -> bpy.types.Modifier:
+def get_lattice_modifier_of_object(obj, lattice) -> 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:
@@ -547,8 +575,8 @@ def get_lattice_modifier_of_object(obj, lattice) -> bpy.types.Modifier:
 
 
 def add_objects_to_lattice(
-        hook: bpy.types.Object,
+        hook: Object,
-        objects: List[bpy.types.Object]):
+        objects: List[Object]):
     lattice_ob = hook['Lattice']
 
     # Check for existing
@@ -568,7 +596,7 @@ def add_objects_to_lattice(
         simple_driver(m, 'strength', hook, '["Tweak Lattice"]')
 
 
-def remove_all_objects_from_lattice(hook: bpy.types.Object) -> List[bpy.types.Object]:
+def remove_all_objects_from_lattice(hook: Object) -> List[Object]:
     lattice_ob = hook['Lattice']
     objs = []
     ob_count = 0
@@ -599,7 +627,13 @@ def remove_objects_from_lattice(hook, objects):
 
 
 classes = [
-    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_Create,
+    TWEAKLAT_OT_Delete,
+    TWEAKLAT_OT_Falloff,
+    TWEAKLAT_OT_Add_Objects,
+    TWEAKLAT_OT_Remove_Selected_Objects,
+    TWEAKLAT_OT_Remove_Object,
+    TWEAKLAT_PT_Main
 ]
 
 
@@ -608,8 +642,8 @@ def register():
     for c in classes:
         register_class(c)
 
-    bpy.types.Scene.tweak_lattice_parent_ob = PointerProperty(
+    Scene.tweak_lattice_parent_ob = PointerProperty(
-        type=bpy.types.Object, name="Parent")
+        type=Object, name="Parent")
 
 
 def unregister():
@@ -617,4 +651,4 @@ def unregister():
     for c in reversed(classes):
         unregister_class(c)
 
-    del bpy.types.Scene.tweak_lattice_parent_ob
+    del Scene.tweak_lattice_parent_ob