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# Copyright 2011, Ryan Inch
20import bpy
21
22from .internals import (
23    layer_collections,
24    qcd_slots,
25    expanded,
26    expand_history,
27    rto_history,
28    copy_buffer,
29    swap_buffer,
30    update_property_group,
31    get_move_selection,
32)
33
34mode_converter = {
35    'EDIT_MESH': 'EDIT',
36    'EDIT_CURVE': 'EDIT',
37    'EDIT_SURFACE': 'EDIT',
38    'EDIT_TEXT': 'EDIT',
39    'EDIT_ARMATURE': 'EDIT',
40    'EDIT_METABALL': 'EDIT',
41    'EDIT_LATTICE': 'EDIT',
42    'POSE': 'POSE',
43    'SCULPT': 'SCULPT',
44    'PAINT_WEIGHT': 'WEIGHT_PAINT',
45    'PAINT_VERTEX': 'VERTEX_PAINT',
46    'PAINT_TEXTURE': 'TEXTURE_PAINT',
47    'PARTICLE': 'PARTICLE_EDIT',
48    'OBJECT': 'OBJECT',
49    'PAINT_GPENCIL': 'PAINT_GPENCIL',
50    'EDIT_GPENCIL': 'EDIT_GPENCIL',
51    'SCULPT_GPENCIL': 'SCULPT_GPENCIL',
52    'WEIGHT_GPENCIL': 'WEIGHT_GPENCIL',
53    'VERTEX_GPENCIL': 'VERTEX_GPENCIL',
54    }
55
56
57rto_path = {
58    "exclude": "exclude",
59    "select": "collection.hide_select",
60    "hide": "hide_viewport",
61    "disable": "collection.hide_viewport",
62    "render": "collection.hide_render",
63    "holdout": "holdout",
64    "indirect": "indirect_only",
65    }
66
67set_off_on = {
68    "exclude": {
69        "off": True,
70        "on": False
71        },
72    "select": {
73        "off": True,
74        "on": False
75        },
76    "hide": {
77        "off": True,
78        "on": False
79        },
80    "disable": {
81        "off": True,
82        "on": False
83        },
84    "render": {
85        "off": True,
86        "on": False
87        },
88    "holdout": {
89        "off": False,
90        "on": True
91        },
92    "indirect": {
93        "off": False,
94        "on": True
95        }
96    }
97
98get_off_on = {
99    False: {
100        "exclude": "on",
101        "select": "on",
102        "hide": "on",
103        "disable": "on",
104        "render": "on",
105        "holdout": "off",
106        "indirect": "off",
107        },
108
109    True: {
110        "exclude": "off",
111        "select": "off",
112        "hide": "off",
113        "disable": "off",
114        "render": "off",
115        "holdout": "on",
116        "indirect": "on",
117        }
118    }
119
120
121def get_rto(layer_collection, rto):
122    if rto in ["exclude", "hide", "holdout", "indirect"]:
123        return getattr(layer_collection, rto_path[rto])
124
125    else:
126        collection = getattr(layer_collection, "collection")
127        return getattr(collection, rto_path[rto].split(".")[1])
128
129
130def set_rto(layer_collection, rto, value):
131    if rto in ["exclude", "hide", "holdout", "indirect"]:
132        setattr(layer_collection, rto_path[rto], value)
133
134    else:
135        collection = getattr(layer_collection, "collection")
136        setattr(collection, rto_path[rto].split(".")[1], value)
137
138
139def apply_to_children(parent, apply_function):
140    # works for both Collections & LayerCollections
141    child_lists = [parent.children]
142
143    while child_lists:
144        new_child_lists = []
145
146        for child_list in child_lists:
147            for child in child_list:
148                apply_function(child)
149
150                if child.children:
151                    new_child_lists.append(child.children)
152
153        child_lists = new_child_lists
154
155
156def isolate_rto(cls, self, view_layer, rto, *, children=False):
157    off = set_off_on[rto]["off"]
158    on = set_off_on[rto]["on"]
159
160    laycol_ptr = layer_collections[self.name]["ptr"]
161    target = rto_history[rto][view_layer]["target"]
162    history = rto_history[rto][view_layer]["history"]
163
164    # get active collections
165    active_layer_collections = [x["ptr"] for x in layer_collections.values()
166                                if get_rto(x["ptr"], rto) == on]
167
168    # check if previous state should be restored
169    if cls.isolated and self.name == target:
170        # restore previous state
171        for x, item in enumerate(layer_collections.values()):
172            set_rto(item["ptr"], rto, history[x])
173
174        # reset target and history
175        del rto_history[rto][view_layer]
176
177        cls.isolated = False
178
179    # check if all RTOs should be activated
180    elif (len(active_layer_collections) == 1 and
181          active_layer_collections[0].name == self.name):
182        # activate all collections
183        for item in layer_collections.values():
184            set_rto(item["ptr"], rto, on)
185
186        # reset target and history
187        del rto_history[rto][view_layer]
188
189        cls.isolated = False
190
191    else:
192        # isolate collection
193
194        rto_history[rto][view_layer]["target"] = self.name
195
196        # reset history
197        history.clear()
198
199        # save state
200        for item in layer_collections.values():
201            history.append(get_rto(item["ptr"], rto))
202
203        child_states = {}
204        if children:
205            # get child states
206            def get_child_states(layer_collection):
207                child_states[layer_collection.name] = get_rto(layer_collection, rto)
208
209            apply_to_children(laycol_ptr, get_child_states)
210
211        # isolate collection
212        for item in layer_collections.values():
213            if item["name"] != laycol_ptr.name:
214                set_rto(item["ptr"], rto, off)
215
216        set_rto(laycol_ptr, rto, on)
217
218        if rto not in ["exclude", "holdout", "indirect"]:
219            # activate all parents
220            laycol = layer_collections[self.name]
221            while laycol["id"] != 0:
222                set_rto(laycol["ptr"], rto, on)
223                laycol = laycol["parent"]
224
225            if children:
226                # restore child states
227                def restore_child_states(layer_collection):
228                    set_rto(layer_collection, rto, child_states[layer_collection.name])
229
230                apply_to_children(laycol_ptr, restore_child_states)
231
232        else:
233            if children:
234                # restore child states
235                def restore_child_states(layer_collection):
236                    set_rto(layer_collection, rto, child_states[layer_collection.name])
237
238                apply_to_children(laycol_ptr, restore_child_states)
239
240            elif rto == "exclude":
241                # deactivate all children
242                def deactivate_all_children(layer_collection):
243                    set_rto(layer_collection, rto, True)
244
245                apply_to_children(laycol_ptr, deactivate_all_children)
246
247        cls.isolated = True
248
249
250def toggle_children(self, view_layer, rto):
251    laycol_ptr = layer_collections[self.name]["ptr"]
252    # clear rto history
253    del rto_history[rto][view_layer]
254    rto_history[rto+"_all"].pop(view_layer, None)
255
256    # toggle rto state
257    state = not get_rto(laycol_ptr, rto)
258    set_rto(laycol_ptr, rto, state)
259
260    def set_state(layer_collection):
261        set_rto(layer_collection, rto, state)
262
263    apply_to_children(laycol_ptr, set_state)
264
265
266def activate_all_rtos(view_layer, rto):
267    off = set_off_on[rto]["off"]
268    on = set_off_on[rto]["on"]
269
270    history = rto_history[rto+"_all"][view_layer]
271
272    # if not activated, activate all
273    if len(history) == 0:
274        keep_history = False
275
276        for item in reversed(list(layer_collections.values())):
277            if get_rto(item["ptr"], rto) == off:
278                keep_history = True
279
280            history.append(get_rto(item["ptr"], rto))
281
282            set_rto(item["ptr"], rto, on)
283
284        if not keep_history:
285            history.clear()
286
287        history.reverse()
288
289    else:
290        for x, item in enumerate(layer_collections.values()):
291            set_rto(item["ptr"], rto, history[x])
292
293        # clear rto history
294        del rto_history[rto+"_all"][view_layer]
295
296
297def invert_rtos(view_layer, rto):
298    if rto == "exclude":
299        orig_values = []
300
301        for item in layer_collections.values():
302            orig_values.append(get_rto(item["ptr"], rto))
303
304        for x, item in enumerate(layer_collections.values()):
305            set_rto(item["ptr"], rto, not orig_values[x])
306
307    else:
308        for item in layer_collections.values():
309            set_rto(item["ptr"], rto, not get_rto(item["ptr"], rto))
310
311    # clear rto history
312    rto_history[rto].pop(view_layer, None)
313
314
315def copy_rtos(view_layer, rto):
316    if not copy_buffer["RTO"]:
317        # copy
318        copy_buffer["RTO"] = rto
319        for laycol in layer_collections.values():
320            copy_buffer["values"].append(get_off_on[
321                                            get_rto(laycol["ptr"], rto)
322                                            ][
323                                            rto
324                                            ]
325                                        )
326
327    else:
328        # paste
329        for x, laycol in enumerate(layer_collections.values()):
330            set_rto(laycol["ptr"],
331                    rto,
332                    set_off_on[rto][
333                        copy_buffer["values"][x]
334                        ]
335                    )
336
337        # clear rto history
338        rto_history[rto].pop(view_layer, None)
339        del rto_history[rto+"_all"][view_layer]
340
341        # clear copy buffer
342        copy_buffer["RTO"] = ""
343        copy_buffer["values"].clear()
344
345
346def swap_rtos(view_layer, rto):
347    if not swap_buffer["A"]["values"]:
348        # get A
349        swap_buffer["A"]["RTO"] = rto
350        for laycol in layer_collections.values():
351            swap_buffer["A"]["values"].append(get_off_on[
352                                                get_rto(laycol["ptr"], rto)
353                                                ][
354                                                rto
355                                                ]
356                                            )
357
358    else:
359        # get B
360        swap_buffer["B"]["RTO"] = rto
361        for laycol in layer_collections.values():
362            swap_buffer["B"]["values"].append(get_off_on[
363                                                get_rto(laycol["ptr"], rto)
364                                                ][
365                                                rto
366                                                ]
367                                            )
368
369        # swap A with B
370        for x, laycol in enumerate(layer_collections.values()):
371            set_rto(laycol["ptr"], swap_buffer["A"]["RTO"],
372                    set_off_on[
373                        swap_buffer["A"]["RTO"]
374                        ][
375                        swap_buffer["B"]["values"][x]
376                        ]
377                    )
378
379            set_rto(laycol["ptr"], swap_buffer["B"]["RTO"],
380                    set_off_on[
381                        swap_buffer["B"]["RTO"]
382                        ][
383                        swap_buffer["A"]["values"][x]
384                        ]
385                    )
386
387
388        # clear rto history
389        swap_a = swap_buffer["A"]["RTO"]
390        swap_b = swap_buffer["B"]["RTO"]
391
392        rto_history[swap_a].pop(view_layer, None)
393        rto_history[swap_a+"_all"].pop(view_layer, None)
394        rto_history[swap_b].pop(view_layer, None)
395        rto_history[swap_b+"_all"].pop(view_layer, None)
396
397        # clear swap buffer
398        swap_buffer["A"]["RTO"] = ""
399        swap_buffer["A"]["values"].clear()
400        swap_buffer["B"]["RTO"] = ""
401        swap_buffer["B"]["values"].clear()
402
403
404def clear_copy(rto):
405    if copy_buffer["RTO"] == rto:
406        copy_buffer["RTO"] = ""
407        copy_buffer["values"].clear()
408
409
410def clear_swap(rto):
411    if swap_buffer["A"]["RTO"] == rto:
412        swap_buffer["A"]["RTO"] = ""
413        swap_buffer["A"]["values"].clear()
414        swap_buffer["B"]["RTO"] = ""
415        swap_buffer["B"]["values"].clear()
416
417
418def link_child_collections_to_parent(laycol, collection, parent_collection):
419    # store view layer RTOs for all children of the to be deleted collection
420    child_states = {}
421    def get_child_states(layer_collection):
422        child_states[layer_collection.name] = (layer_collection.exclude,
423                                               layer_collection.hide_viewport,
424                                               layer_collection.holdout,
425                                               layer_collection.indirect_only)
426
427    apply_to_children(laycol["ptr"], get_child_states)
428
429    # link any subcollections of the to be deleted collection to it's parent
430    for subcollection in collection.children:
431        if not subcollection.name in parent_collection.children:
432            parent_collection.children.link(subcollection)
433
434    # apply the stored view layer RTOs to the newly linked collections and their
435    # children
436    def restore_child_states(layer_collection):
437        state = child_states.get(layer_collection.name)
438
439        if state:
440            layer_collection.exclude = state[0]
441            layer_collection.hide_viewport = state[1]
442            layer_collection.holdout = state[2]
443            layer_collection.indirect_only = state[3]
444
445    apply_to_children(laycol["parent"]["ptr"], restore_child_states)
446
447
448def remove_collection(laycol, collection, context):
449    # get selected row
450    cm = context.scene.collection_manager
451    selected_row_name = cm.cm_list_collection[cm.cm_list_index].name
452
453    # delete collection
454    bpy.data.collections.remove(collection)
455
456    # update references
457    expanded.discard(laycol["name"])
458
459    if expand_history["target"] == laycol["name"]:
460        expand_history["target"] = ""
461
462    if laycol["name"] in expand_history["history"]:
463        expand_history["history"].remove(laycol["name"])
464
465    if qcd_slots.contains(name=laycol["name"]):
466        qcd_slots.del_slot(name=laycol["name"])
467
468    if laycol["name"] in qcd_slots.overrides:
469        qcd_slots.overrides.remove(laycol["name"])
470
471    # reset history
472    for rto in rto_history.values():
473        rto.clear()
474
475    # update tree view
476    update_property_group(context)
477
478    # update selected row
479    laycol = layer_collections.get(selected_row_name, None)
480    if laycol:
481        cm.cm_list_index = laycol["row_index"]
482
483    elif len(cm.cm_list_collection) <= cm.cm_list_index:
484        cm.cm_list_index =  len(cm.cm_list_collection) - 1
485
486        if cm.cm_list_index > -1:
487            name = cm.cm_list_collection[cm.cm_list_index].name
488            laycol = layer_collections[name]
489            while not laycol["visible"]:
490                laycol = laycol["parent"]
491
492            cm.cm_list_index = laycol["row_index"]
493
494
495def select_collection_objects(is_master_collection, collection_name, replace, nested, selection_state=None):
496    if bpy.context.mode != 'OBJECT':
497        return
498
499    if is_master_collection:
500        target_collection = bpy.context.view_layer.layer_collection.collection
501
502    else:
503        laycol = layer_collections[collection_name]
504        target_collection = laycol["ptr"].collection
505
506    if replace:
507        bpy.ops.object.select_all(action='DESELECT')
508
509    if selection_state == None:
510        selection_state = get_move_selection().isdisjoint(target_collection.objects)
511
512    def select_objects(collection):
513            for obj in collection.objects:
514                try:
515                    obj.select_set(selection_state)
516                except RuntimeError:
517                    pass
518
519    select_objects(target_collection)
520
521    if nested:
522        apply_to_children(target_collection, select_objects)
523
524def set_exclude_state(target_layer_collection, state):
525    # get current child exclusion state
526    child_exclusion = []
527
528    def get_child_exclusion(layer_collection):
529        child_exclusion.append([layer_collection, layer_collection.exclude])
530
531    apply_to_children(target_layer_collection, get_child_exclusion)
532
533
534    # set exclusion
535    target_layer_collection.exclude = state
536
537
538    # set correct state for all children
539    for laycol in child_exclusion:
540        laycol[0].exclude = laycol[1]
541