source/pie_selection.py

@@ -3,8 +3,9 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import bpy
-from bpy.types import Menu
+from bpy.types import Menu, Operator
 from .hotkeys import register_hotkey
+from bpy.props import BoolProperty
 
 
 class PIE_MT_object_selection(Menu):
@@ -19,15 +20,15 @@ class PIE_MT_object_selection(Menu):
         # 6 - RIGHT
         pie.operator("object.select_all", text="Select All", icon='OUTLINER_OB_POINTCLOUD').action='SELECT'
         # 2 - BOTTOM
-        pie.operator("object.select_grouped", text="Direct Children", icon='CON_CHILDOF').type='CHILDREN'
+        pie.operator("object.select_children_of_active", text=f"Direct Children", icon='CON_CHILDOF')
         # 8 - TOP
-        pie.operator("object.select_grouped", text="Parent", icon='FILE_PARENT').type='PARENT'
+        pie.operator("object.select_parent_object", text="Parent", icon='FILE_PARENT')
         # 7 - TOP - LEFT
         pie.operator("object.select_all", text="Invert Selection", icon='CLIPUV_DEHLT').action='INVERT'
         # 9 - TOP - RIGHT
-        pie.operator("object.select_grouped", text="Siblings", icon='PARTICLES').type='SIBLINGS'
+        pie.operator("object.select_siblings_of_active", text="Siblings", icon='PARTICLES')
         # 1 - BOTTOM - LEFT
-        pie.operator("object.select_grouped", text="All Children", icon='OUTLINER').type='CHILDREN_RECURSIVE'
+        pie.operator("object.select_children_of_active", text=f"All Children", icon='OUTLINER').recursive=True
         # 3 - BOTTOM - RIGHT
         pie.menu("PIE_MT_object_selection_more", text="Select Menu", icon='THREE_DOTS')
 
@@ -70,15 +71,6 @@ class PIE_MT_object_selection_more(Menu):
         layout.operator("view3d.select_box", text="Box Select", icon='SELECT_SET')
         layout.operator_menu_enum("view3d.select_lasso", "mode", icon='MOD_CURVE')
 
-        # Parent/Children selection extend ops.
-        layout.separator()
-        props = layout.operator("object.select_hierarchy", text="Add Parent to Selection", icon='FILE_PARENT')
-        props.extend = True
-        props.direction = 'PARENT'
-        props = layout.operator("object.select_hierarchy", text="Add Children to Selection", icon='CON_CHILDOF')
-        props.extend = True
-        props.direction = 'CHILD'
-
         # Select based on the active object.
         layout.separator()
         layout.operator_menu_enum("object.select_grouped", "type", text="Select Grouped", icon='OUTLINER_COLLECTION')
@@ -131,11 +123,161 @@ class PIE_MT_mesh_selection_more(Menu):
         layout.template_node_operator_asset_menu_items(catalog_path="Select")
 
 
+def get_selected_parents(context) -> list:
+    """Return objects which are selected or active, and have children.
+    Active object is always first."""
+    objects = context.selected_objects or [context.active_object]
+    parents = [o for o in objects if o.children]
+    if context.active_object and context.active_object in parents:
+        idx = parents.index(context.active_object)
+        if idx > -1:
+            parents.pop(idx)
+        parents.insert(0, context.active_object)
+    return parents
+
+
+def deselect_all_objects(context):
+    for obj in context.selected_objects:
+        obj.select_set(False)
+
+
+class ObjectSelectOperatorMixin:
+    extend_selection: BoolProperty(
+        name="Extend Selection",
+        description="Bones that are already selected will remain selected",
+    )
+
+    def invoke(self, context, event):
+        if event.shift:
+            self.extend_selection = True
+        else:
+            self.extend_selection = False
+
+        return self.execute(context)
+
+    @classmethod
+    def poll(cls, context):
+        if context.mode != 'OBJECT':
+            cls.poll_message_set("Must be in Object Mode.")
+            return False
+        return True
+
+    def execute(self, context):
+        if not self.extend_selection:
+            deselect_all_objects(context)
+
+        return {'FINISHED'}
+
+
+class OBJECT_OT_select_children_of_active(Operator, ObjectSelectOperatorMixin):
+    """Select the children of the active object"""
+
+    bl_idname = "object.select_children_of_active"
+    bl_label = "Select Children"
+    bl_options = {'REGISTER', 'UNDO'}
+
+    recursive: BoolProperty(default=False, options={'SKIP_SAVE'})
+
+    @classmethod
+    def poll(cls, context):
+        if not super().poll(context):
+            return False
+        obj = context.active_object
+        if not (obj and obj.children):
+            cls.poll_message_set("Active object has no children.")
+            return False
+        return True
+
+    @classmethod
+    def description(cls, context, props):
+        recursively = " recursively" if props.recursive else ""
+        return f"Select children of active object{recursively}." + "\n\nShift: Extend current selection"
+
+    def execute(self, context):
+        super().execute(context)
+
+        counter = 0
+        children = context.active_object.children
+        if self.recursive:
+            children = context.active_object.children_recursive
+        for child in children:
+            if not child.select_get():
+                counter += 1
+                child.select_set(True)
+
+        self.report({'INFO'}, f"Selected {counter} objects.")
+        return {'FINISHED'}
+
+
+class OBJECT_OT_select_siblings_of_active(Operator, ObjectSelectOperatorMixin):
+    """Select siblings of active object.\n\nShift: Extend current selection"""
+
+    bl_idname = "object.select_siblings_of_active"
+    bl_label = "Select Siblings"
+    bl_options = {'REGISTER', 'UNDO'}
+
+    @classmethod
+    def poll(cls, context):
+        if not super().poll(context):
+            return False
+        obj = context.active_object
+        if not (obj and obj.parent):
+            cls.poll_message_set("Active object has no parent.")
+            return False
+        if not len(obj.parent.children) > 1:
+            cls.poll_message_set("Active object has no siblings.")
+            return False
+        return True
+
+    def execute(self, context):
+        super().execute(context)
+
+        counter = 0
+        for child in context.active_object.parent.children:
+            if not child.select_get():
+                counter += 1
+                child.select_set(True)
+
+        self.report({'INFO'}, f"Selected {counter} objects.")
+        return {'FINISHED'}
+
+
+class OBJECT_OT_select_parent_object(Operator, ObjectSelectOperatorMixin):
+    """Select parent of the active object.\n\nShift: Extend current selection"""
+
+    bl_idname = "object.select_parent_object"
+    bl_label = "Select Parent Object"
+    bl_options = {'REGISTER', 'UNDO'}
+
+    @classmethod
+    def poll(cls, context):
+        if not super().poll(context):
+            return False
+        obj = context.active_object
+        if not (obj and obj.parent):
+            cls.poll_message_set("Active object has no parent.")
+            return False
+        return True
+
+    def execute(self, context):
+        super().execute(context)
+
+        active_obj = context.active_object
+
+        active_obj.parent.select_set(True)
+        context.view_layer.objects.active = active_obj.parent
+
+        return {'FINISHED'}
+
+
 registry = [
     PIE_MT_object_selection,
     PIE_MT_mesh_selection,
     PIE_MT_object_selection_more,
     PIE_MT_mesh_selection_more,
+    OBJECT_OT_select_children_of_active,
+    OBJECT_OT_select_siblings_of_active,
+    OBJECT_OT_select_parent_object,
 ]