From 257c5e099be7191920069a4163845543abe7d6d8 Mon Sep 17 00:00:00 2001
From: Jakub Trllo <jakub.trllo@gmail.com>
Date: Fri, 18 Feb 2022 10:47:30 +0100
Subject: [PATCH] fix right click in representation widget

---
 openpype/tools/loader/model.py   | 14 ++++++++++++
 openpype/tools/loader/widgets.py | 38 +++++++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/openpype/tools/loader/model.py b/openpype/tools/loader/model.py
index f34eb157ab..10b22d0e17 100644
--- a/openpype/tools/loader/model.py
+++ b/openpype/tools/loader/model.py
@@ -1068,6 +1068,7 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
         self._icons = lib.get_repre_icons()
         self._icons["repre"] = qtawesome.icon("fa.file-o",
                                               color=style.colors.default)
+        self._items_by_id = {}
 
     def set_version_ids(self, version_ids):
         self.version_ids = version_ids
@@ -1076,6 +1077,9 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
     def data(self, index, role):
         item = index.internalPointer()
 
+        if role == ITEM_ID_ROLE:
+            return item["id"]
+
         if role == self.IdRole:
             return item.get("_id")
 
@@ -1149,8 +1153,11 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
             if len(self.version_ids) > 1:
                 group = repre_groups.get(doc["name"])
                 if not group:
+
                     group_item = Item()
+                    item_id = str(uuid4())
                     group_item.update({
+                        "id": item_id,
                         "_id": doc["_id"],
                         "name": doc["name"],
                         "isMerged": True,
@@ -1161,6 +1168,7 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
                             color=style.colors.default
                         )
                     })
+                    self._items_by_id[item_id] = group_item
                     self.add_child(group_item, None)
                     repre_groups[doc["name"]] = group_item
                     repre_groups_items[doc["name"]] = 0
@@ -1173,7 +1181,9 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
             active_site_icon = self._icons.get(self.active_provider)
             remote_site_icon = self._icons.get(self.remote_provider)
 
+            item_id = str(uuid4())
             data = {
+                "id": item_id,
                 "_id": doc["_id"],
                 "name": doc["name"],
                 "subset": doc["context"]["subset"],
@@ -1192,6 +1202,7 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
 
             item = Item()
             item.update(data)
+            self._items_by_id[item_id] = item
 
             current_progress = {
                 'active_site_progress': progress[self.active_site],
@@ -1215,6 +1226,9 @@ class RepresentationModel(TreeModel, BaseRepresentationModel):
         self.endResetModel()
         self.refreshed.emit(False)
 
+    def get_item_by_id(self, item_id):
+        return self._items_by_id.get(item_id)
+
     def refresh(self):
         docs = []
         session_project = self.dbcon.Session['AVALON_PROJECT']
diff --git a/openpype/tools/loader/widgets.py b/openpype/tools/loader/widgets.py
index f60341af2f..f145756cc5 100644
--- a/openpype/tools/loader/widgets.py
+++ b/openpype/tools/loader/widgets.py
@@ -1283,6 +1283,40 @@ class RepresentationWidget(QtWidgets.QWidget):
             }
         return repre_context_by_id
 
+    def get_selected_items(self):
+        selection_model = self.tree_view.selectionModel()
+        indexes = selection_model.selectedIndexes()
+
+        item_ids = set()
+        for index in indexes:
+            item_id = index.data(ITEM_ID_ROLE)
+            if item_id is not None:
+                item_ids.add(item_id)
+
+        output = []
+        for item_id in item_ids:
+            item = self.model.get_item_by_id(item_id)
+            if item is not None:
+                output.append(item)
+        return output
+
+    def get_selected_repre_items(self):
+        output = []
+        items = collections.deque(self.get_selected_items())
+
+        item_ids = set()
+        while items:
+            item = items.popleft()
+            if item.get("isGroup") or item.get("isMerged"):
+                for child in item.children():
+                    items.appendleft(child)
+            else:
+                item_id = item["id"]
+                if item_id not in item_ids:
+                    item_ids.add(item_id)
+                    output.append(item)
+        return output
+
     def on_context_menu(self, point):
         """Shows menu with loader actions on Right-click.
 
@@ -1301,10 +1335,8 @@ class RepresentationWidget(QtWidgets.QWidget):
         selection = self.tree_view.selectionModel()
         rows = selection.selectedRows(column=0)
 
-        items = self.get_selected_subsets()
-
+        items = self.get_selected_repre_items()
         selected_side = self._get_selected_side(point_index, rows)
-
         # Get all representation->loader combinations available for the
         # index under the cursor, so we can list the user the options.
         available_loaders = api.discover(api.Loader)