refactoring caption code

Author: choldgraf

docs/demo/index.rst

@@ -16,6 +16,7 @@ the left sidebar to see how various elements look on this theme.
     lists_tables
     markdown
     example_pandas
+    index2
 
 .. toctree::
     :maxdepth: 3

docs/demo/index2.rst

@@ -0,0 +1,8 @@
+A second index
+==============
+
+An index to show off level nesting.
+
+.. toctree::
+
+    page2

docs/demo/page2.rst

@@ -0,0 +1,4 @@
+A second page
+=============
+
+A demo page to see how levels nest within each other in the theme.

pydata_sphinx_theme/__init__.py

@@ -20,28 +20,27 @@
 
 class MyTocTree(TocTree):
     def get_toctree_for_subpage(
-        self, pagename, builder, collapse=True, maxdepth=-1, **kwargs
+        self, toctree_page, current_page, builder, collapse=True, maxdepth=-1, **kwargs
     ):
         """Return the global TOC nodetree."""
-        if pagename in ["genindex", "search"]:
+        if current_page in ["genindex", "search"]:
             return
-        doctree = self.env.get_doctree(pagename)
-        toctrees = []  # type: List[Element]
-        if "includehidden" not in kwargs:
-            kwargs["includehidden"] = True
-        if "maxdepth" not in kwargs:
-            kwargs["maxdepth"] = 0
-        kwargs["collapse"] = collapse
+        doctree = self.env.get_doctree(
+            toctree_page
+        )  # current_page should be the "first active page"
+        toctrees = []
         for toctreenode in doctree.traverse(addnodes.toctree):
-            toctree = self.resolve(pagename, builder, toctreenode, prune=True, **kwargs)
+            toctree = self.resolve(
+                current_page, builder, toctreenode, prune=True, **kwargs
+            )
             if toctree:
                 toctrees.append(toctree)
         if not toctrees:
             return None
-        result = toctrees[0]
+        toctree_subpage = toctrees[0]
         for toctree in toctrees[1:]:
-            result.extend(toctree.children)
-        return result
+            toctree_subpage.extend(toctree.children)
+        return toctree_subpage
 
 
 def add_toctree_functions(app, pagename, templatename, context, doctree):
@@ -77,31 +76,48 @@ def get_nav_object(maxdepth=None, collapse=True, subpage_caption=False, **kwargs
             return []
 
         if subpage_caption:
-            if pagename not in [app.env.config.master_doc, "genindex", "search"]:
-                def is_first_active_page(node):
-                    return isinstance(node, nodes.bullet_list) and node.attributes.get("iscurrent")
-
-                active_first_page = list(toctree.traverse(is_first_active_page))[0]
-                # A path to the active TOC item's first page, relative to the current page
-                first_page_path = list(active_first_page.traverse(nodes.reference))[0].attributes.get("refuri")
-                if first_page_path == "":
-                    # First TOC item's first page *is* the active page
-                    first_page_path = Path(pagename).name
-                else:
-                    first_page_path = Path(first_page_path).with_suffix("")
-                rel_first_page_path = str(Path(pagename).parent.joinpath(first_page_path))
-
-                # We only wish to show a single page's descendants, so we'll keep their captions
-                subpage_toctree = toc.get_toctree_for_subpage(
-                    rel_first_page_path, app.builder, collapse=collapse, maxdepth=maxdepth, **kwargs
+
+            def is_level1_page(node):
+                return isinstance(
+                    node, nodes.list_item
+                ) and "toctree-l1" in node.attributes.get("classes")
+
+            # Grab all of the first-level pages from the index toctree
+            toctree_root = toc.get_toctree_for(
+                app.env.config.master_doc,
+                app.builder,
+                collapse=collapse,
+                maxdepth=maxdepth,
+                **kwargs,
+            )
+            first_pages = toctree_root.traverse(is_level1_page)
+            # For each page in the top toctree, access the page-specific toctree
+            for first_page in first_pages:
+                # First reference will be for the top page
+                first_page_ref = list(first_page.traverse(nodes.reference))[0]
+                first_page_path = first_page_ref.attributes.get("refuri").rsplit(
+                        ".html", 1
+                )[0]
+
+                # Grab the toctree for this page, with links relative from current page
+                first_page_toctree = toc.get_toctree_for_subpage(
+                    first_page_path,
+                    pagename,
+                    app.builder,
+                    collapse=collapse,
+                    maxdepth=maxdepth - 1,
+                    **kwargs,
                 )
-                if subpage_toctree is not None:
+
+                if first_page_toctree is not None:
                     # Find the current page in the top-level children
-                    for item in toctree.children:
-                        if isinstance(item, nodes.bullet_list) and item.attributes.get("iscurrent", []):
-                            # Append that pages' toctree so we get captions
-                            subpage_list = item.children[0]
-                            subpage_list.children = [subpage_list.children[0]] + subpage_toctree.children        
+                    for page in toctree.traverse(is_level1_page):
+                        toc_page_ref = list(page.traverse(nodes.reference))[0]
+                        page_uri = toc_page_ref.attributes.get("refuri").rsplit(
+                            ".html", 1
+                        )[0]
+                        if page_uri == first_page_path:
+                            page.children = first_page_toctree.children
 
         # toctree has this structure
         #   <caption>