[pitivi] timeline: Fix selection post grouping/ungrouping

[Alexandru Băluț <alexbalut src gnome org>]

commit e8a9753cb153bd541d2336476d318a5d449b9f9d
Author: Alexandru Băluț <alexandru balut gmail com>
Date:   Sat Feb 5 22:41:17 2022 +0100

    timeline: Fix selection post grouping/ungrouping
    
    Previously, if you selected an a/v clip and ungroup it, then reselect
    the two resulting clips and re-group them, the resulting clip would
    appear with the video selected and audio deselected.
    
    Previously, if you ungrouped a group/clip, the selection would reset.
    
    Now the selection persists grouping and ungrouping.

 pitivi/timeline/timeline.py     | 36 ++++++++++---------
 pitivi/utils/timeline.py        | 78 ++++++++++++++++++++++++-----------------
 tests/common.py                 |  2 --
 tests/test_timeline_timeline.py | 65 ++++++++++++++++++----------------
 tests/test_undo_timeline.py     |  5 ++-
 tests/test_utils_timeline.py    | 12 +++----
 6 files changed, 108 insertions(+), 90 deletions(-)
---
diff --git a/pitivi/timeline/timeline.py b/pitivi/timeline/timeline.py
index e868d0b44..1361b903a 100644
--- a/pitivi/timeline/timeline.py
+++ b/pitivi/timeline/timeline.py
@@ -1752,14 +1752,14 @@ class TimelineContainer(Gtk.Grid, Zoomable, Loggable):
 
         # Clips actions.
         self.delete_action = Gio.SimpleAction.new("delete-selected-clips", None)
-        self.delete_action.connect("activate", self._delete_selected)
+        self.delete_action.connect("activate", self._delete_selected_cb)
         group.add_action(self.delete_action)
         self.app.shortcuts.add("timeline.delete-selected-clips", ["Delete"],
                                self.delete_action,
                                _("Delete selected clips"))
 
         self.delete_and_shift_action = Gio.SimpleAction.new("delete-selected-clips-and-shift", None)
-        self.delete_and_shift_action.connect("activate", self._delete_selected_and_shift)
+        self.delete_and_shift_action.connect("activate", self._delete_selected_and_shift_cb)
         group.add_action(self.delete_and_shift_action)
         self.app.shortcuts.add("timeline.delete-selected-clips-and-shift", ["<Shift>Delete"],
                                self.delete_and_shift_action,
@@ -1978,7 +1978,7 @@ class TimelineContainer(Gtk.Grid, Zoomable, Loggable):
         self.timeline.update_position()
         return False
 
-    def _delete_selected(self, unused_action, unused_parameter):
+    def _delete_selected_cb(self, unused_action, unused_parameter):
         if self.ges_timeline:
             with Previewer.manager.paused():
                 with self.app.action_log.started("delete clip",
@@ -1992,7 +1992,7 @@ class TimelineContainer(Gtk.Grid, Zoomable, Loggable):
 
             self.timeline.selection.set_selection([], SELECT)
 
-    def _delete_selected_and_shift(self, unused_action, unused_parameter):
+    def _delete_selected_and_shift_cb(self, unused_action, unused_parameter):
         if self.ges_timeline:
             with self.app.action_log.started("delete clip and shift",
                                              
finalizing_action=CommitTimelineFinalizingAction(self._project.pipeline),
@@ -2036,35 +2036,37 @@ class TimelineContainer(Gtk.Grid, Zoomable, Loggable):
 
     def _ungroup_selected_cb(self, unused_action, unused_parameter):
         if not self.ges_timeline:
-            self.info("No ges_timeline set yet!")
             return
 
+        toplevels = self.timeline.selection.toplevels()
+        containers = set()
         with self.app.action_log.started("ungroup",
                                          
finalizing_action=CommitTimelineFinalizingAction(self._project.pipeline),
                                          toplevel=True):
-            toplevels = self.timeline.selection.toplevels
             for toplevel in toplevels:
+                # Normally we don't ungroup Clips unless they are
+                # selected by themselves.
                 if isinstance(toplevel, GES.Group) or len(toplevels) == 1:
-                    toplevel.ungroup(recursive=False)
+                    containers |= set(toplevel.ungroup(recursive=False))
 
-        self.timeline.selection.set_selection([], SELECT)
+        clips = list(self.timeline.selection.get_clips_of(containers))
+        self.timeline.selection.set_selection(clips, SELECT)
 
     def _group_selected_cb(self, unused_action, unused_parameter):
         if not self.ges_timeline:
-            self.info("No timeline set yet?")
+            return
+
+        toplevels = self.timeline.selection.toplevels()
+        if not toplevels:
             return
 
         with self.app.action_log.started("group",
                                          
finalizing_action=CommitTimelineFinalizingAction(self._project.pipeline),
                                          toplevel=True):
-            toplevels = self.timeline.selection.toplevels
-            if toplevels:
-                GES.Container.group(list(toplevels))
-
-            # timeline.selection doesn't change during grouping,
-            # we need to manually update group actions.
-            self.timeline.selection.set_can_group_ungroup()
-            self.update_actions()
+            container = GES.Container.group(list(toplevels))
+
+        clips = self.timeline.selection.get_clips_of([container])
+        self.timeline.selection.set_selection(clips, SELECT)
 
     def __copy_clips_cb(self, unused_action, unused_parameter):
         group = self.timeline.selection.group()
diff --git a/pitivi/utils/timeline.py b/pitivi/utils/timeline.py
index 976e9a9f4..4eb91d000 100644
--- a/pitivi/utils/timeline.py
+++ b/pitivi/utils/timeline.py
@@ -15,7 +15,7 @@
 #
 # You should have received a copy of the GNU Lesser General Public
 # License along with this program; if not, see <http://www.gnu.org/licenses/>.
-from typing import List
+from typing import Iterable
 from typing import Set
 
 from gi.repository import GES
@@ -74,7 +74,8 @@ class Selection(GObject.Object, Loggable):
     """Manages a set of clips representing a selection.
 
     Signals:
-        selection-changed: The contents of the selection changed.
+        selection-changed: The contents of the selection or the grouping of
+            the selected clips changed.
     """
 
     __gsignals__ = {
@@ -88,60 +89,64 @@ class Selection(GObject.Object, Loggable):
         self.can_group = False
         self.can_ungroup = False
 
-    def set_selection(self, clips: List[GES.Clip], mode: int):
+    def set_selection(self, clips: Iterable[GES.Clip], mode: int):
         """Updates the current selection.
 
         Args:
-            clips (List[GES.Clip]): Timeline clips to update the selection with.
+            clips (Iterable[GES.Clip]): Timeline clips to update the selection with.
             mode (SELECT or UNSELECT or SELECT_ADD): The type of update to
                 apply. The selection will be:
                 - `SELECT` : set to the provided selection.
                 - `UNSELECT` : the same minus the provided selection.
                 - `SELECT_ADD` : extended with the provided selection.
         """
+        if not isinstance(clips, Set):
+            clips = set(clips)
         self.debug("Updating selection %s mode %s", clips, mode)
         if mode == SELECT_ADD:
-            selection = self._clips | set(clips)
+            selection = self._clips | clips
         elif mode == UNSELECT:
-            selection = self._clips - set(clips)
+            selection = self._clips - clips
         else:
-            selection = set(clips)
+            selection = clips
 
-        if self._clips == selection:
-            # Nothing changed. This can happen for example when the user clicks
-            # the selected clip, then the clip remains selected.
-            return
+        selection_changed = self._clips != selection
+        if selection_changed:
+            old_selection = self._clips
+            self._clips = selection
 
-        old_selection = self._clips
-        self._clips = selection
+            for obj, selected in self.__get_selection_changes(old_selection):
+                obj.selected.selected = selected
+                if obj.ui:
+                    from pitivi.utils.ui import set_state_flags_recurse
+                    set_state_flags_recurse(obj.ui, Gtk.StateFlags.SELECTED, are_set=selected)
 
-        for obj, selected in self.__get_selection_changes(old_selection):
-            obj.selected.selected = selected
-            if obj.ui:
-                from pitivi.utils.ui import set_state_flags_recurse
-                set_state_flags_recurse(obj.ui, Gtk.StateFlags.SELECTED, are_set=selected)
+                for element in obj.get_children(False):
+                    if isinstance(obj, (GES.BaseEffect, GES.TextOverlay)):
+                        continue
+                    element.selected.selected = selected
 
-            for element in obj.get_children(False):
-                if isinstance(obj, (GES.BaseEffect, GES.TextOverlay)):
-                    continue
-                element.selected.selected = selected
+        # Always check, maybe the clips have been grouped or ungrouped.
+        grouping_changed = self._set_can_group_ungroup()
 
-        self.set_can_group_ungroup()
+        if selection_changed or grouping_changed:
+            self.emit("selection-changed")
 
-        self.emit("selection-changed")
-
-    def set_can_group_ungroup(self):
-        can_ungroup = False
-        toplevels = self.toplevels
-        self.can_group = len(toplevels) > 1
+    def _set_can_group_ungroup(self) -> bool:
+        toplevels = self.toplevels()
+        can_group = len(toplevels) > 1
         # If we allow grouping, we disallow ungrouping.
-        if not self.can_group:
+        can_ungroup = False
+        if not can_group:
             for toplevel in toplevels:
                 if (isinstance(toplevel, GES.Group) or
                         (isinstance(toplevel, GES.Container) and len(toplevel.get_children(False)) > 1)):
                     can_ungroup = True
                     break
-        self.can_ungroup = can_ungroup and not self.can_group
+        changed = (self.can_group, self.can_ungroup) != (can_group, can_ungroup)
+        self.can_group = can_group
+        self.can_ungroup = can_ungroup
+        return changed
 
     def __get_selection_changes(self, old_selection):
         for obj in old_selection - self._clips:
@@ -183,7 +188,6 @@ class Selection(GObject.Object, Loggable):
                 return clip
         return None
 
-    @property
     def toplevels(self):
         """Returns the toplevel elements of the selection."""
         toplevels = set()
@@ -232,6 +236,16 @@ class Selection(GObject.Object, Loggable):
     def __iter__(self):
         return iter(self._clips)
 
+    @staticmethod
+    def get_clips_of(containers: Iterable[GES.Container]) -> Iterable[GES.Clip]:
+        for container in containers:
+            if isinstance(container, GES.Clip):
+                yield container
+            elif isinstance(container, GES.Container):
+                for child in container.get_children(recursive=True):
+                    if isinstance(child, GES.Clip):
+                        yield child
+
 
 class EditingContext(GObject.Object, Loggable):
     """Encapsulates interactive editing.
diff --git a/tests/common.py b/tests/common.py
index 5ea2028ca..3129b7c0b 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -432,8 +432,6 @@ class TestCase(unittest.TestCase, Loggable):
 
     def click_clip(self, ges_clip: GES.Clip, expect_selected: bool, ctrl_key: bool = False):
         """Clicks the specified clip."""
-        self.assert_clip_selected(ges_clip, not expect_selected)
-
         timeline = ges_clip.ui.timeline
         original_control_mask = timeline.get_parent().control_mask
         if ctrl_key:
diff --git a/tests/test_timeline_timeline.py b/tests/test_timeline_timeline.py
index dfcda10c4..d59dee1fa 100644
--- a/tests/test_timeline_timeline.py
+++ b/tests/test_timeline_timeline.py
@@ -241,45 +241,44 @@ class TestGrouping(common.TestCase):
     def test_can_group_ungroup(self):
         timeline_container = common.create_timeline_container()
         timeline = timeline_container.timeline
+        # Timeline empty.
         self.__check_can_group_ungroup(timeline_container, False, False)
+
         ges_clip, = self.add_clips_simple(timeline, 1)
+        # No clip selected.
+        self.__check_can_group_ungroup(timeline_container, False, False)
+
         self.click_clip(ges_clip, expect_selected=True)
+        # An audio-video clip is selected.
         self.__check_can_group_ungroup(timeline_container, False, True)
 
         timeline_container.ungroup_action.emit("activate", None)
-        self.__check_can_group_ungroup(timeline_container, False, False)
+        # The resulting audio clip and video clip should both be selected.
+        self.__check_can_group_ungroup(timeline_container, True, False)
 
         layer, = timeline.ges_timeline.get_layers()
         ges_clip0, ges_clip1 = layer.get_clips()
-        self.click_clip(ges_clip0, expect_selected=True)
+        self.click_clip(ges_clip0, expect_selected=False, ctrl_key=True)
+        # Only one of the clips remains selected.
         self.__check_can_group_ungroup(timeline_container, False, False)
 
-        # Press <ctrl> so selecting in ADD mode
-        event = mock.Mock()
-        event.keyval = Gdk.KEY_Control_L
-        timeline_container.do_key_press_event(event)
-
-        self.click_clip(ges_clip1, expect_selected=True)
+        self.click_clip(ges_clip0, expect_selected=True)
+        self.click_clip(ges_clip1, expect_selected=True, ctrl_key=True)
+        # Both clip are selected.
         self.__check_can_group_ungroup(timeline_container, True, False)
 
         timeline_container.group_action.emit("activate", None)
+        # The resulting audio-video clip should be selected.
         self.__check_can_group_ungroup(timeline_container, False, True)
 
     def group_clips(self, timeline_container, clips):
         timeline = timeline_container.timeline
         timeline.app.settings.leftClickAlsoSeeks = False
 
-        # Press <ctrl> so selecting in ADD mode
-        event = mock.Mock()
-        event.keyval = Gdk.KEY_Control_L
-        timeline_container.do_key_press_event(event)
-        timeline.get_clicked_layer_and_pos = mock.Mock()
-        timeline.get_clicked_layer_and_pos.return_value = (None, None)
-
         # Select the 2 clips
         for clip in clips:
             self.assertIsNone(clip.get_parent())
-            self.click_clip(clip, expect_selected=True)
+            self.click_clip(clip, expect_selected=True, ctrl_key=True)
 
         timeline_container.group_action.emit("activate", None)
 
@@ -296,6 +295,8 @@ class TestGrouping(common.TestCase):
         group = clips[0].get_parent()
         self.assertEqual(len(group.get_children(False)), len(clips))
 
+        self.assertEqual(len(timeline.selection), len(clips))
+
     def test_group(self):
         timeline_container = common.create_timeline_container()
         timeline = timeline_container.timeline
@@ -308,17 +309,20 @@ class TestGrouping(common.TestCase):
         timeline = timeline_container.timeline
         clips = self.add_clips_simple(timeline, num_clips)
         self.group_clips(timeline_container, clips)
-        layer = timeline.ges_timeline.get_layers()[0]
+        layer, = timeline.ges_timeline.get_layers()
         clips = layer.get_clips()
         self.assertEqual(len(clips), num_clips)
 
-        # Deselect one grouped clip clips
-        self.click_clip(clips[0], expect_selected=False)
-
-        # Make sure all the clips have been deselected
+        # Deselect one of the clips in the group.
+        self.click_clip(clips[0], expect_selected=False, ctrl_key=True)
         for clip in clips:
             self.assert_clip_selected(clip, expect_selected=False)
 
+        # Select one of the clips in the group.
+        self.click_clip(clips[0], expect_selected=True)
+        for clip in clips:
+            self.assert_clip_selected(clip, expect_selected=True)
+
     def test_group_ungroup(self):
         num_clips = 2
         timeline_container = common.create_timeline_container()
@@ -326,7 +330,8 @@ class TestGrouping(common.TestCase):
         clips = self.add_clips_simple(timeline, num_clips)
         self.group_clips(timeline_container, clips)
 
-        self.assertEqual(len(timeline.selection), num_clips)
+        # Selecting a clip selects all the clips in its group.
+        self.click_clip(clips[0], expect_selected=True)
 
         timeline_container.ungroup_action.emit("activate", None)
         layer = timeline.ges_timeline.get_layers()[0]
@@ -372,14 +377,6 @@ class TestGrouping(common.TestCase):
         # Group the two selected clips
         timeline_container.group_action.emit("activate", None)
 
-        # Deselect the second clip, notice both clips have been deselected
-        self.click_clip(clips[1], expect_selected=False, ctrl_key=True)
-        self.assert_clip_selected(clips[0], expect_selected=False)
-
-        # Select the second clip, notice both clips have been selected
-        self.click_clip(clips[1], expect_selected=True)
-        self.assert_clip_selected(clips[0], expect_selected=True)
-
         # Deselect the first clip, notice both clips have been deselected
         self.click_clip(clips[0], expect_selected=False, ctrl_key=True)
         self.assert_clip_selected(clips[1], expect_selected=False)
@@ -388,6 +385,14 @@ class TestGrouping(common.TestCase):
         self.click_clip(clips[0], expect_selected=True)
         self.assert_clip_selected(clips[1], expect_selected=True)
 
+        # Deselect the second clip, notice both clips have been deselected
+        self.click_clip(clips[1], expect_selected=False, ctrl_key=True)
+        self.assert_clip_selected(clips[0], expect_selected=False)
+
+        # Select the second clip, notice both clips have been selected
+        self.click_clip(clips[1], expect_selected=True)
+        self.assert_clip_selected(clips[0], expect_selected=True)
+
     def test_ungroup_clip(self):
         timeline_container = common.create_timeline_container()
         timeline = timeline_container.timeline
diff --git a/tests/test_undo_timeline.py b/tests/test_undo_timeline.py
index a06c34379..5a68ead5e 100644
--- a/tests/test_undo_timeline.py
+++ b/tests/test_undo_timeline.py
@@ -154,14 +154,13 @@ class TestTimelineObserver(common.TestCase):
 
         self.assertTrue(self.layer.add_clip(clip1))
         self.assertTrue(self.layer.add_clip(clip2))
-        # The selection does not care about GES.Groups, only about GES.Clips.
-        self.timeline_container.timeline.selection.select([clip1, clip2])
         self.assertEqual(clip1.props.timeline, self.layer.get_timeline())
         self.assertEqual(clip2.props.timeline, self.layer.get_timeline())
 
+        self.timeline_container.timeline.selection.select([clip1, clip2])
         self.timeline_container.group_action.activate(None)
-        self.assertTrue(isinstance(clip1.get_parent(), GES.Group))
         self.assertEqual(clip1.get_parent(), clip2.get_parent())
+        self.assertTrue(isinstance(clip1.get_parent(), GES.Group), type(clip1.get_parent()))
 
         self.timeline_container.ungroup_action.activate(None)
         self.assertIsNone(clip1.get_parent())
diff --git a/tests/test_utils_timeline.py b/tests/test_utils_timeline.py
index 042ce6fae..99dcfa75c 100644
--- a/tests/test_utils_timeline.py
+++ b/tests/test_utils_timeline.py
@@ -107,27 +107,27 @@ class TestSelection(common.TestCase):
         selection = Selection()
 
         selection.set_selection([clip1, clip2, clip3, clip4], SELECT)
-        self.assertSetEqual(selection.toplevels, {clip1, clip2, clip3, clip4})
+        self.assertSetEqual(selection.toplevels(), {clip1, clip2, clip3, clip4})
 
         group1 = GES.Container.group([clip1, clip2])
         group1.props.serialize = True
-        self.assertSetEqual(selection.toplevels, {group1, clip3, clip4})
+        self.assertSetEqual(selection.toplevels(), {group1, clip3, clip4})
 
         group2 = GES.Container.group([group1, clip3])
         group2.props.serialize = True
-        self.assertSetEqual(selection.toplevels, {group2, clip4})
+        self.assertSetEqual(selection.toplevels(), {group2, clip4})
 
         group1.props.serialize = True
         group1.props.serialize = False
-        self.assertSetEqual(selection.toplevels, {group2, clip4})
+        self.assertSetEqual(selection.toplevels(), {group2, clip4})
 
         group1.props.serialize = False
         group2.props.serialize = False
-        self.assertSetEqual(selection.toplevels, {clip1, clip2, clip3, clip4})
+        self.assertSetEqual(selection.toplevels(), {clip1, clip2, clip3, clip4})
 
         group1.props.serialize = True
         group2.props.serialize = False
-        self.assertSetEqual(selection.toplevels, {group1, clip3, clip4})
+        self.assertSetEqual(selection.toplevels(), {group1, clip3, clip4})
 
 
 class TestEditingContext(common.TestCase):