E-book viewer: Allow expanding/collapsing all items in the Table of C…

…ontents at a particular level by right clicking on one item of that level and choosing the option to expand/collapse

src/calibre/gui2/viewer/toc.py

@@ -103,6 +103,16 @@ def expand_tree(self, index):
                 break
             self.expand_tree(child)
 
+    def collapse_at_level(self, index):
+        item = self.model().itemFromIndex(index)
+        for x in self.model().items_at_depth(item.depth):
+            self.collapse(self.model().indexFromItem(x))
+
+    def expand_at_level(self, index):
+        item = self.model().itemFromIndex(index)
+        for x in self.model().items_at_depth(item.depth):
+            self.expand(self.model().indexFromItem(x))
+
     def context_menu(self, pos):
         index = self.indexAt(pos)
         m = QMenu(self)
@@ -112,6 +122,10 @@ def context_menu(self, pos):
         m.addAction(_('Expand all items'), self.expandAll)
         m.addAction(_('Collapse all items'), self.collapseAll)
         m.addSeparator()
+        if index.isValid():
+            m.addAction(_('Expand all items at the level of {}').format(index.data()), partial(self.expand_at_level, index))
+            m.addAction(_('Collapse all items at the level of {}').format(index.data()), partial(self.collapse_at_level, index))
+        m.addSeparator()
         m.addAction(_('Copy table of contents to clipboard'), self.copy_to_clipboard)
         m.exec_(self.mapToGlobal(pos))
 
@@ -161,7 +175,7 @@ def do_search(self, text):
 
 class TOCItem(QStandardItem):
 
-    def __init__(self, toc, depth, all_items, normal_font, emphasis_font, parent=None):
+    def __init__(self, toc, depth, all_items, normal_font, emphasis_font, depths, parent=None):
         text = toc.get('title') or ''
         self.href = (toc.get('dest') or '')
         if toc.get('frag'):
@@ -174,8 +188,10 @@ def __init__(self, toc, depth, all_items, normal_font, emphasis_font, parent=Non
         QStandardItem.__init__(self, text)
         all_items.append(self)
         self.normal_font, self.emphasis_font = normal_font, emphasis_font
-        for t in toc['children']:
-            self.appendRow(TOCItem(t, depth+1, all_items, normal_font, emphasis_font, parent=self))
+        if toc['children']:
+            depths.add(depth + 1)
+            for t in toc['children']:
+                self.appendRow(TOCItem(t, depth+1, all_items, normal_font, emphasis_font, depths, parent=self))
         self.setFlags(Qt.ItemFlag.ItemIsEnabled)
         self.is_current_search_result = False
         self.depth = depth
@@ -223,9 +239,11 @@ def __init__(self, toc=None):
         normal_font = QApplication.instance().font()
         emphasis_font = QFont(normal_font)
         emphasis_font.setBold(True), emphasis_font.setItalic(True)
+        self.depths = {0}
         if toc:
             for t in toc['children']:
-                self.appendRow(TOCItem(t, 0, depth_first, normal_font, emphasis_font))
+                self.appendRow(TOCItem(t, 0, depth_first, normal_font, emphasis_font, self.depths))
+        self.depths = tuple(sorted(self.depths))
         self.node_id_map = {x.node_id: x for x in self.all_items}
 
     def find_items(self, query):
@@ -234,6 +252,11 @@ def find_items(self, query):
             if not query or (text and primary_contains(query, text)):
                 yield item
 
+    def items_at_depth(self, depth):
+        for item in self.all_items:
+            if item.depth == depth:
+                yield item
+
     def node_id_for_text(self, query):
         for item in self.find_items(query):
             return item.node_id