anim_selection_sets.py

 #  ***** BEGIN GPL LICENSE BLOCK *****
 #
 #  This program is free software; you can redistribute it and/or
 #  modify it under the terms of the GNU General Public License
 #  as published by the Free Software Foundation; either version 2
 #  of the License, or (at your option) any later version.
 #
 #  This program is distributed in the hope that it will be useful,
 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 #  GNU General Public License for more details.
 #
 #  You should have received a copy of the GNU General Public License
 #  along with this program; if not, see <http://www.gnu.org/licenses/>
 #  and write to the Free Software Foundation, Inc., 51 Franklin Street,
 #  Fifth Floor, Boston, MA  02110-1301, USA..
 #
 #  The Original Code is Copyright (C) 2012 Blender Foundation ###
 #  All rights reserved.
 #
 #
 #  The Original Code is: all of this file.
 #
 #  Contributor(s): Dan Eicher.
 #
 #  ***** END GPL LICENSE BLOCK *****

# <pep8 compliant>

import string
import bpy

bl_info = {
  "name": "Selection Set",
  "author": "Dan Eicher",
  "version": (0, 1, 1),
  "blender": (2, 65, 4),
  "location": "Properties > Object data (Armature) > Selection Sets",
  "description": "Selection Sets to select groups of posebones",
  "warning": "Proxy armatures need to export sets and "
    "run generated script on re-opening file",
  "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/Py/"
    "Scripts/Animation/SelectionSets",
  "tracker_url": "https://developer.blender.org/T31492",
  "category": "Animation"
}


script_template = '''# generated by ANIM_OT_selection_set_export -- abandon all hope, ye who hand edit!

import bpy

def selection_set_import():
    arm = bpy.data.armatures['${name}']

    set_list = [${set_list}
               ]

    for name, bones in set_list:
        if arm.selection_sets.find(name) == -1:
            sel_set = arm.selection_sets.add()
            sel_set.name = name

            for bone in bones:
                set_bone = sel_set.bones.add()
                set_bone.name = bone


if __name__ == "__main__":
    selection_set_import()
'''


def generic_poll(context):
    return (context.mode == 'POSE' and context.object and context.object.type == 'ARMATURE')


def active_selection_set_update_func(self, context):
    idx = self.active_selection_set
    if idx < -1 or idx >= len(self.selection_sets):
        self.active_selection_set = -1
        raise IndexError('Armature.active_selection_set: out of range')


class SelectionSetBone(bpy.types.PropertyGroup):
    name = bpy.props.StringProperty()


class SelectionSet(bpy.types.PropertyGroup):
    name = bpy.props.StringProperty(options={'LIBRARY_EDITABLE'})
    bones = bpy.props.CollectionProperty(type=SelectionSetBone)


class ANIM_OT_selection_set_add(bpy.types.Operator):
    """Add a new selection set"""
    bl_idname = "anim.selection_set_add"
    bl_label = "Selection Set Add"

    @classmethod
    def poll(cls, context):
        return generic_poll(context)

    def execute(self, context):
        arm = context.active_object.data

        tmp_name = name = 'Set'
        name_sub = 1
        while arm.selection_sets.find(name) != -1:
            name = tmp_name + ' ' + str(name_sub)
            name_sub += 1

        sel_set = arm.selection_sets.add()
        sel_set.name = name

        arm.active_selection_set = arm.selection_sets.find(name)

        return {'FINISHED'}


class ANIM_OT_selection_set_remove(bpy.types.Operator):
    """Remove the active selection set"""
    bl_idname = "anim.selection_set_remove"
    bl_label = "Selection Set Remove"

    @classmethod
    def poll(cls, context):
        arm = context.active_object.data
        return (generic_poll(context) and arm.active_selection_set != -1)

    def execute(self, context):
        arm = context.active_object.data
        active_index = arm.active_selection_set

        arm.selection_sets.remove(active_index)

        if active_index >= len(arm.selection_sets):
            arm.active_selection_set = len(arm.selection_sets) - 1

        return {'FINISHED'}


class ANIM_OT_selection_set_assign(bpy.types.Operator):
    """Add selected bones to the active selection set"""
    bl_idname = "anim.selection_set_assign"
    bl_label = "Selection Set Assign"

    @classmethod
    def poll(cls, context):
        arm = context.active_object.data
        return (generic_poll(context) and arm.active_selection_set != -1)

    def execute(self, context):
        arm = context.active_object.data
        sel_set = arm.selection_sets[arm.active_selection_set]
        bones = [bone for bone in arm.bones if bone.select]

        for bone in bones:
            if sel_set.bones.find(bone.name) == -1:
                set_bone = sel_set.bones.add()
                set_bone.name = bone.name

        return {'FINISHED'}


class ANIM_OT_selection_set_unassign(bpy.types.Operator):
    """Remove selected bones from the active selection set"""
    bl_idname = "anim.selection_set_unassign"
    bl_label = "Selection Set Unassign"

    @classmethod
    def poll(cls, context):
        arm = context.active_object.data
        return (generic_poll(context) and arm.active_selection_set != -1)

    def execute(self, context):
        arm = context.active_object.data
        sel_set = arm.selection_sets[arm.active_selection_set]
        bones = [bone for bone in arm.bones if bone.select]

        for bone in bones:
            bone_index = sel_set.bones.find(bone.name)
            if bone_index != -1:
                sel_set.bones.remove(bone_index)

        return {'FINISHED'}


class ANIM_OT_selection_set_select(bpy.types.Operator):
    """Select bones in selection set"""
    bl_idname = "anim.selection_set_select"
    bl_label = "Selection Set Select Bones"

    @classmethod
    def poll(cls, context):
        arm = context.active_object.data
        return (generic_poll(context) and arm.active_selection_set != -1)

    def execute(self, context):
        arm = context.active_object.data
        sel_set = arm.selection_sets[arm.active_selection_set]

        for bone in sel_set.bones:
            try:
                arm.bones[bone.name].select = True
            except:
                bone_index = sel_set.bones.find(bone.name)
                sel_set.bones.remove(bone_index)

        return {'FINISHED'}


class ANIM_OT_selection_set_deselect(bpy.types.Operator):
    """Deselect bones in selection set"""
    bl_idname = "anim.selection_set_deselect"
    bl_label = "Selection Set Deselect Bones"

    @classmethod
    def poll(cls, context):
        arm = context.active_object.data
        return (generic_poll(context) and arm.active_selection_set != -1)

    def execute(self, context):
        arm = context.active_object.data
        sel_set = arm.selection_sets[arm.active_selection_set]

        for bone in sel_set.bones:
            try:
                arm.bones[bone.name].select = False
            except:
                bone_index = sel_set.bones.find(bone.name)
                sel_set.bones.remove(bone_index)

        return {'FINISHED'}


class ANIM_OT_selection_set_export(bpy.types.Operator):
    """Export selection set data to a python script"""
    bl_idname = "anim.selection_set_export"
    bl_label = "Selection Set Export"

    @classmethod
    def poll(cls, context):
        return generic_poll(context)

    def execute(self, context):
        arm = context.active_object.data
        set_script = string.Template(script_template)
        set_list = ""

        for sel_set in arm.selection_sets:
            set_bones = ""
            for bone in sel_set.bones:
                set_bones += "'" + bone.name + "',"
            set_list += "\n                ('{name}', [{bones}]),".format(name=sel_set.name, bones=set_bones)

        try:
            script_file = bpy.data.texts['{arm.name}SelectionSetImport.py'.format(arm=arm)]
        except:
            script_file = bpy.data.texts.new('{arm.name}SelectionSetImport.py'.format(arm=arm))

        script_file.clear()
        script_file.write(set_script.substitute(name=arm.name, set_list=set_list))

        return {'FINISHED'}


class DATA_PT_bone_sets(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "data"
    bl_label = "Selection Sets"

    @classmethod
    def poll(cls, context):
        return (context.object and context.object.type == 'ARMATURE' and context.object.pose)

    def draw(self, context):
        layout = self.layout
        ob = context.object
        arm = ob.data
        sel_set = None

        if arm.active_selection_set != -1:
            try:
                sel_set = arm.selection_sets[arm.active_selection_set]
            except:
                pass

        row = layout.row()

        row.template_list("UI_UL_list", "armature_selection_sets", arm, "selection_sets", arm, "active_selection_set",
                          rows=(5 if len(arm.selection_sets) else 2))

        col = row.column(align=True)
        col.operator("anim.selection_set_add", icon='ZOOMIN', text="")
        col.operator("anim.selection_set_remove", icon='ZOOMOUT', text="")

        if sel_set:
            col = layout.column()
            col.prop(sel_set, "name", text="Name")

        row = layout.row()

        sub = row.row(align=True)
        sub.operator("anim.selection_set_assign", text="Assign")
        sub.operator("anim.selection_set_unassign", text="Remove")

        sub = row.row(align=True)
        sub.operator("anim.selection_set_select", text="Select")
        sub.operator("anim.selection_set_deselect", text="Deselect")

        row = layout.row()
        row.operator("anim.selection_set_export", text="Export Selection Sets")


def register():
    bpy.utils.register_class(SelectionSetBone)
    bpy.utils.register_class(SelectionSet)
    bpy.utils.register_class(ANIM_OT_selection_set_add)
    bpy.utils.register_class(ANIM_OT_selection_set_remove)
    bpy.utils.register_class(ANIM_OT_selection_set_assign)
    bpy.utils.register_class(ANIM_OT_selection_set_unassign)
    bpy.utils.register_class(ANIM_OT_selection_set_select)
    bpy.utils.register_class(ANIM_OT_selection_set_deselect)
    bpy.utils.register_class(ANIM_OT_selection_set_export)
    bpy.utils.register_class(DATA_PT_bone_sets)

    bpy.types.Armature.selection_sets = bpy.props.CollectionProperty(type=SelectionSet,
                                            name="Selection Sets",
                                            description="Collection of bones to be selected")

    bpy.types.Armature.active_selection_set = bpy.props.IntProperty(name="Active Selection Set",
                                                                    default=-1,
                                                                    options={'LIBRARY_EDITABLE'},
                                                                    update=active_selection_set_update_func)


def unregister():
    del bpy.types.Armature.selection_sets
    del bpy.types.Armature.active_selection_set

    bpy.utils.unregister_class(SelectionSet)
    bpy.utils.unregister_class(SelectionSetBone)
    bpy.utils.unregister_class(ANIM_OT_selection_set_add)
    bpy.utils.unregister_class(ANIM_OT_selection_set_remove)
    bpy.utils.unregister_class(ANIM_OT_selection_set_assign)
    bpy.utils.unregister_class(ANIM_OT_selection_set_unassign)
    bpy.utils.unregister_class(ANIM_OT_selection_set_select)
    bpy.utils.unregister_class(ANIM_OT_selection_set_deselect)
    bpy.utils.unregister_class(ANIM_OT_selection_set_export)
    bpy.utils.unregister_class(DATA_PT_bone_sets)

if __name__ == "__main__":
    register()