object_collections.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 import bpy
  20 from bpy.types import (
  21         Operator,
  22         Panel,
  23         Menu,
  24         )
  25
  26
  27 from bpy.props import (
  28         EnumProperty,
  29         IntProperty,
  30         )
  31
  32
  33 bl_info = {
  34     "name": "Collections",
  35     "author": "Dalai Felinto",
  36     "version": (1, 0),
  37     "blender": (2, 80, 0),
  38     "description": "Panel to set/unset object collections",
  39     "warning": "",
  40     "wiki_url": "",
  41     "tracker_url": "",
  42     "category": "Object"}
  43
  44
  45 # #####################################################################################
  46 # Operators
  47 # #####################################################################################
  48
  49 class OBJECT_OT_collection_add(Operator):
  50     """Add an object to a new collection"""
  51     bl_idname = "object.collection_add"
  52     bl_label = "Add to New Collection"
  53
  54     def execute(self, context):
  55         scene = context.scene
  56         collection = scene.master_collection.collections.new()
  57         collection.objects.link(context.object)
  58         return {'FINISHED'}
  59
  60
  61 class OBJECT_OT_collection_remove(Operator):
  62     """Remove the active object from this collection"""
  63     bl_idname = "object.collection_remove"
  64     bl_label = "Remove from Collection"
  65
  66     def execute(self, context):
  67         collection = context.scene_collection
  68         collection.objects.unlink(context.object)
  69         return {'FINISHED'}
  70
  71
  72 def get_collection_from_id_recursive(collection, collection_id, current_id):
  73     """Return len of collection and the collection if it was a match"""
  74     if collection_id == current_id:
  75         return collection, 0
  76
  77     current_id += 1
  78
  79     for collection_nested in collection.collections:
  80         matched_collection, current_id = get_collection_from_id_recursive(
  81                                                  collection_nested,
  82                                                  collection_id,
  83                                                  current_id)
  84         if matched_collection is not None:
  85             return matched_collection, 0
  86
  87     return None, current_id
  88
  89
  90 def get_collection_from_id(scene, collection_id):
  91     master_collection = scene.master_collection
  92     return get_collection_from_id_recursive(master_collection, collection_id, 0)[0]
  93
  94
  95 def collection_items_recursive(path, collection, items, current_id, object_name):
  96     name = collection.name
  97     current_id += 1
  98
  99     if object_name not in collection.objects:
 100         items.append((str(current_id), path + name, ""))
 101         path += name + " / "
 102
 103     for collection_nested in collection.collections:
 104         current_id = collection_items_recursive(path, collection_nested, items, current_id, object_name)
 105     return current_id
 106
 107
 108 def collection_items(self, context):
 109     items = []
 110
 111     master_collection = context.scene.master_collection
 112     object_name = context.object.name
 113
 114     if object_name not in master_collection.objects:
 115         items.append(('0', "Master Collection", "", 'COLLAPSEMENU', 0))
 116
 117     current_id = 0
 118     for collection in master_collection.collections:
 119         current_id = collection_items_recursive("", collection, items, current_id, object_name)
 120
 121     return items
 122
 123
 124 class OBJECT_OT_collection_link(Operator):
 125     """Add an object to an existing collection"""
 126     bl_idname = "object.collection_link"
 127     bl_label = "Link to Collection"
 128
 129     collection_index = IntProperty(
 130             name = "Collection Index",
 131             default = -1,
 132             options = {'SKIP_SAVE'},
 133             )
 134
 135     type = EnumProperty(
 136             name = "",
 137             description = "Dynamic enum for collections",
 138             items=collection_items,
 139             )
 140
 141     def execute(self, context):
 142         if self.collection_index == -1:
 143             self.collection_index = int(self.type)
 144
 145         collection = get_collection_from_id(context.scene, self.collection_index)
 146
 147         if collection is None:
 148             # It should never ever happen!
 149             self.report({'ERROR'}, "Unexpected error: collection {0} is invalid".format(
 150                     self.collection_index))
 151             return {'CANCELLED'}
 152
 153         collection.objects.link(context.object)
 154         return {'FINISHED'}
 155
 156     def invoke(self, context, events):
 157         if self.collection_index != -1:
 158             return self.execute(context)
 159
 160         wm = context.window_manager
 161         wm.invoke_search_popup(self)
 162         return {'FINISHED'}
 163
 164
 165 def find_collection_parent(collection, collection_parent):
 166     for collection_nested in collection_parent.collections:
 167         if collection_nested == collection:
 168             return collection_parent
 169
 170         found_collection = find_collection_parent(collection, collection_nested)
 171         if found_collection:
 172             return found_collection
 173     return None
 174
 175
 176 class OBJECT_OT_collection_unlink(Operator):
 177     """Unlink the collection from all objects"""
 178     bl_idname = "object.collection_unlink"
 179     bl_label = "Unlink Collection"
 180
 181     def execute(self, context):
 182         collection = context.scene_collection
 183         master_collection = context.scene.master_collection
 184
 185         collection_parent = find_collection_parent(collection, master_collection)
 186         if collection_parent is None:
 187             self.report({'ERROR'}, "Cannot find {0}'s parent".format(collection.name))
 188             return {'CANCELLED'}
 189
 190         collection_parent.collections.remove(collection)
 191         return {'CANCELLED'}
 192
 193
 194 def select_collection_objects(collection):
 195     for ob in collection.objects:
 196         ob.select_set(True)
 197
 198     for collection_nested in collection.collections:
 199         select_collection_objects(collection_nested)
 200
 201
 202 class OBJECT_OT_collection_select(Operator):
 203     """Select all objects in collection"""
 204     bl_idname = "object.collection_select"
 205     bl_label = "Select Collection"
 206
 207     def execute(self, context):
 208         collection = context.scene_collection
 209         select_collection_objects(collection)
 210         return {'FINISHED'}
 211
 212
 213 # #####################################################################################
 214 # Interface
 215 # #####################################################################################
 216
 217 class COLLECTION_MT_specials(Menu):
 218     bl_label = "Collection Specials"
 219
 220     def draw(self, context):
 221         layout = self.layout
 222
 223         col = layout.column()
 224         col.active = context.scene_collection != context.scene.master_collection
 225         col.operator("object.collection_unlink", icon='X', text="Unlink Collection")
 226
 227         layout.operator("object.collection_select", text="Select Collection")
 228
 229
 230 def all_collections_get(context):
 231     """Iterator over all scene collections
 232     """
 233     def all_collections_recursive_get(collection_parent, collections):
 234         collections.append(collection_parent)
 235         for collection_nested in collection_parent.collections:
 236             all_collections_recursive_get(collection_nested, collections)
 237
 238     scene = context.scene
 239     master_collection = scene.master_collection
 240
 241     collections = []
 242
 243     all_collections_recursive_get(master_collection, collections)
 244
 245     return collections
 246
 247
 248 class OBJECT_PT_collections(Panel):
 249     bl_space_type = 'PROPERTIES'
 250     bl_region_type = 'WINDOW'
 251     bl_context = "object"
 252     bl_label = "Collections"
 253
 254     def draw(self, context):
 255         layout = self.layout
 256         row = layout.row(align=True)
 257
 258         obj = context.object
 259         master_collection = bpy.context.scene.master_collection
 260
 261         if master_collection.collections:
 262             row.operator("object.collection_link", text="Add to Collection")
 263         else:
 264             row.operator("object.collection_link", text="Add to Collection").collection_index = 0
 265         row.operator("object.collection_add", text="", icon='ADD')
 266
 267         obj_name = obj.name
 268         for collection in all_collections_get(context):
 269             collection_objects = collection.objects
 270             if obj_name in collection.objects:
 271                 col = layout.column(align=True)
 272                 col.context_pointer_set("scene_collection", collection)
 273
 274                 row = col.box().row()
 275                 if collection == master_collection:
 276                     row.label(text=collection.name)
 277                 else:
 278                     row.prop(collection, "name", text="")
 279                 row.operator("object.collection_remove", text="", icon='X', emboss=False)
 280                 row.menu("COLLECTION_MT_specials", icon='DOWNARROW_HLT', text="")
 281
 282
 283 # #####################################################################################
 284 # Register/Unregister
 285 # #####################################################################################
 286
 287 classes = (
 288     COLLECTION_MT_specials,
 289     OBJECT_PT_collections,
 290     OBJECT_OT_collection_add,
 291     OBJECT_OT_collection_remove,
 292     OBJECT_OT_collection_link,
 293     OBJECT_OT_collection_unlink,
 294     OBJECT_OT_collection_select,
 295 )
 296
 297
 298 def register():
 299     for cls in classes:
 300         bpy.utils.register_class(cls)
 301
 302
 303 def unregister():
 304     for cls in classes:
 305         bpy.utils.unregister_class(cls)
 306
 307
 308 if __name__ == "__main__":
 309     register()