Rewrite transforms, merge post-transforms, fix pep-naming

Also fix bug with `hNode` not being reset in `LatexMasterDocTransforms`

Author: chrisjsewell

.pre-commit-config.yaml

@@ -39,7 +39,10 @@ repos:
     rev: 3.8.4
     hooks:
     - id: flake8
-      additional_dependencies: [flake8-bugbear==21.3.1]
+      additional_dependencies:
+      - flake8-bugbear~=21.3.1
+      - flake8-comprehensions~=3.4.0
+      - pep8-naming~=0.11.1
 
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v0.790

jupyterbook_latex/events.py

@@ -11,10 +11,9 @@
 
 from . import __version__, theme
 from .transforms import (
-    LatexMasterDocTransforms,
+    LatexRootDocPostTransforms,
+    LatexRootDocTransforms,
     MystNbPostTransform,
-    ToctreeTransforms,
-    handleSubSections,
 )
 
 if sys.version_info < (3, 9):
@@ -81,10 +80,10 @@ def setup_latex_transforms(app: Sphinx) -> None:
     if MystNbPostTransform.check_dependency():
         app.add_post_transform(MystNbPostTransform)
 
-    TOC_PATH = Path(app.confdir or app.srcdir).joinpath("_toc.yml")
-    if not os.path.exists(TOC_PATH):
+    toc_path = Path(app.confdir or app.srcdir).joinpath("_toc.yml")
+    if not os.path.exists(toc_path):
         logger.info(
-            "Some features of this exetension will work only with a jupyter-book application"  # noqa: E501
+            "Some features of this extension will work only with a jupyter-book application"  # noqa: E501
         )
         return
 
@@ -99,6 +98,5 @@ def setup_latex_transforms(app: Sphinx) -> None:
         )
 
     app.setup_extension("sphinx.ext.imgconverter")
-    app.add_transform(LatexMasterDocTransforms)
-    app.add_post_transform(ToctreeTransforms)
-    app.add_post_transform(handleSubSections)
+    app.add_transform(LatexRootDocTransforms)
+    app.add_post_transform(LatexRootDocPostTransforms)

jupyterbook_latex/transforms.py

@@ -1,11 +1,12 @@
 from pathlib import Path
-from typing import Any, List
+from typing import Any, Dict, List, Optional, Tuple, Type
 
 import docutils
 import yaml
 from sphinx import addnodes
 from sphinx.application import Sphinx
 from sphinx.builders.latex.nodes import thebibliography
+from sphinx.environment import BuildEnvironment
 from sphinx.transforms import SphinxTransform
 from sphinx.transforms.post_transforms import SphinxPostTransform
 from sphinx.util import logging
@@ -15,27 +16,38 @@
 logger = logging.getLogger(__name__)
 
 
-def replaceWithNode(srcNode, toReplace, copyChildren):
-    node = toReplace("")
-    node["classes"] = srcNode["classes"]
-    if copyChildren:
-        node.children = srcNode.children
-    srcNode.replace_self([node])
+def replace_node_cls(
+    src_node: docutils.nodes.Element,
+    node_cls: Type[docutils.nodes.Element],
+    copy_children: bool,
+) -> None:
+    """Replace the class of a node."""
+    node = node_cls("")
+    node["classes"] = src_node["classes"]
+    if copy_children:
+        node.children = src_node.children
+    src_node.replace_self([node])
 
 
-def depth(node, parentId):
+def get_section_depth(node: docutils.nodes.section, root_ids: Tuple[str, ...]) -> int:
+    """Get the nesting depth of the section."""
     d = 0
-    while node.attributes["ids"][0] != parentId:
+    while tuple(node.attributes["ids"]) != root_ids:
         d += 1
-        node = node.parent
+        node = getattr(node, "parent", None)
+        if not node:
+            break
     return d
 
 
-def find_parent(env, node, parentTag):
-    while 1:
+def find_parent(
+    env: BuildEnvironment, node: docutils.nodes.Element, parent_tag: str
+) -> Optional[docutils.nodes.Element]:
+    """Find the parent node."""
+    while True:
         node = node.parent
         if node is None:
-            return
+            return None
         # parent should be a document in toc
         if (
             "docname" in node.attributes
@@ -44,7 +56,7 @@ def find_parent(env, node, parentTag):
         ):
             return node
 
-    if node.tagname == parentTag:
+    if node.tagname == parent_tag:
         return node
 
     return None
@@ -63,54 +75,55 @@ def is_root_document(document: docutils.nodes.document, app: Sphinx) -> bool:
     return app.project.path2doc(document["source"]) == app.config.master_doc
 
 
-# Transforms and postTransforms
+class LatexRootDocTransforms(SphinxTransform):
+    """Arrange the toctrees and sections in the required structure."""
 
-
-class LatexMasterDocTransforms(SphinxTransform):
     default_priority = 500
 
     def apply(self, **kwargs: Any) -> None:
-        tocwrrapper = []
-
-        def getSectionDepth():
-            levelDict = dict()
-            for count, sect in enumerate(
-                self.document.traverse(docutils.nodes.section)
-            ):
-                if count == 0:
-                    parentSect = sect
-                    parentId = sect.attributes["ids"][0]
-                levelDict[sect.attributes["ids"][0]] = depth(sect, parentId)
-            return levelDict, parentSect
-
-        def alterNodes(sectLevelsDict, parentSect):
-            hNode = None
-            for sect in self.document.traverse(docutils.nodes.section):
-                if parentSect != sect:
-                    for child in sect.children:
-                        if isinstance(child, docutils.nodes.title):
-                            if sectLevelsDict[sect.attributes["ids"][0]] == 1:
-                                hNode = H2Node("")
-                            elif sectLevelsDict[sect.attributes["ids"][0]] == 2:
-                                hNode = H3Node("")
-                            if hNode is not None:
-                                hNode.children = child.children
-                                child = hNode
-                        parentSect.append(child)
-                    replaceWithNode(sect, HiddenCellNode, False)
-
-        if is_root_document(self.document, self.app):
-            # pull the toctree-wrapper and append it later to the topmost document level
-            for node in self.document.traverse(docutils.nodes.compound):
-                if "toctree-wrapper" in node["classes"]:
-                    tocwrrapper.append(node)
-                    replaceWithNode(node, HiddenCellNode, False)
-
-            sectLevelsDict, parentSect = getSectionDepth()
-            alterNodes(sectLevelsDict, parentSect)
-            # append the toctreewrapper to the topmost level of document
-            if tocwrrapper:
-                self.document.extend(tocwrrapper)
+
+        if not is_root_document(self.document, self.app):
+            return
+
+        # pop the toctree-wrappers and append them to the topmost document level
+        for node in self.document.traverse(docutils.nodes.compound):
+            if "toctree-wrapper" in node["classes"]:
+                replace_node_cls(node, HiddenCellNode, False)
+                self.document.append(node)
+
+        # map the section ids to their depth in the tree (h1 == 0)
+        section_levels: Dict[Tuple[str, ...], int] = {}
+        for count, sect in enumerate(self.document.traverse(docutils.nodes.section)):
+            sect_id = tuple(sect.attributes["ids"])
+            if count == 0:
+                top_level_section = sect
+                top_level_id = sect_id
+            section_levels[sect_id] = get_section_depth(sect, top_level_id)
+
+        if not section_levels:
+            return
+
+        # flatten the AST sections under the top-level section
+        for sect in self.document.traverse(docutils.nodes.section):
+            if sect == top_level_section:
+                continue
+            # move section children to the top-level
+            for child in sect.children:
+                # replace H2 and H3 titles with nodes that can be custom rendered
+                if isinstance(child, docutils.nodes.title):
+                    sect_id = tuple(sect.attributes["ids"])
+                    header_node = None
+                    if section_levels[sect_id] == 1:
+                        header_node = H2Node("")
+                    elif section_levels[sect_id] == 2:
+                        header_node = H3Node("")
+                    if header_node is not None:
+                        header_node.children = child.children
+                        child = header_node
+                top_level_section.append(child)
+
+            # remove the section node
+            replace_node_cls(sect, HiddenCellNode, False)
 
 
 class MystNbPostTransform(SphinxPostTransform):
@@ -140,80 +153,76 @@ def apply(self, **kwargs: Any) -> None:
 
         for node in self.document.traverse(CellNode):
             if "tag_hide-cell" in node["classes"]:
-                replaceWithNode(node, HiddenCellNode, True)
+                replace_node_cls(node, HiddenCellNode, True)
             if "tag_hide-input" in node["classes"]:
-                inputNode = node.traverse(CellInputNode)
-                for node in inputNode:
-                    replaceWithNode(node, HiddenCellNode, True)
+                for input_node in node.traverse(CellInputNode):
+                    replace_node_cls(input_node, HiddenCellNode, True)
             if "tag_hide-output" in node["classes"]:
-                outputNode = node.traverse(CellOutputNode)
-                for node in outputNode:
-                    replaceWithNode(node, HiddenCellNode, True)
+                for output_node in node.traverse(CellOutputNode):
+                    replace_node_cls(output_node, HiddenCellNode, True)
 
 
-class handleSubSections(SphinxPostTransform):
-    default_priority = 700
+def check_node_in_part(part, node, app: Sphinx) -> bool:
+    """ """
+    if "chapters" in part:
+        nodefile = node.children[0].attributes["docname"]
+        chapfiles = part["chapters"]
+        for chap in chapfiles:
+            chapname = remove_suffix(list(chap.values())[0], app.config.source_suffix)
+            if nodefile in chapname:
+                return True
+    return False
 
-    def apply(self, **kwargs: Any) -> None:
-        if is_root_document(self.document, self.app):
-            for compound in self.document.traverse(docutils.nodes.compound):
-                if "toctree-wrapper" in compound["classes"]:
-                    nodecopy = compound
-                    node = find_parent(self.app.env, nodecopy, "section")
-                    if node:
-                        replaceWithNode(compound, HiddenCellNode, False)
-                        node.append(nodecopy)
 
+class LatexRootDocPostTransforms(SphinxPostTransform):
+    """Arrange the toctrees and bibliographies into the required structure."""
 
-class ToctreeTransforms(SphinxPostTransform):
-    default_priority = 800
+    default_priority = 700
 
     def apply(self, **kwargs: Any) -> None:
-        def checkNodeIsInPart(part, node):
-            if "chapters" in part:
-                nodefile = node.children[0].attributes["docname"]
-                chapfiles = part["chapters"]
-                for chap in chapfiles:
-                    chapname = remove_suffix(
-                        list(chap.values())[0], self.app.config.source_suffix
-                    )
-                    if nodefile in chapname:
-                        return True
-            return False
+        if not is_root_document(self.document, self.app):
+            return
 
-        if is_root_document(self.document, self.app):
-            TOC_PATH = Path(self.app.confdir or self.app.srcdir).joinpath("_toc.yml")
-            tocfile = yaml.safe_load(TOC_PATH.read_text("utf8"))
-
-            # store bibliography nodes to append it at the end of doc
-            bibNodes = []
-            for bibnode in self.document.traverse(thebibliography):
-                bibNodes.append(bibnode)
-                replaceWithNode(bibnode, HiddenCellNode, False)
-
-            for f in tocfile:
-                if "part" in f:
-                    self.app.config["latex_toplevel_sectioning"] = "part"
-                    partname = f["part"]
-                    compoundParent = docutils.nodes.compound("")
-                    compoundParent["classes"] = "toctree-wrapper"
-                    startOfFile = addnodes.start_of_file("")
-                    startOfFile["docname"] = partname
-                    title = docutils.nodes.title(text=partname)
-                    sectionName = docutils.nodes.section("")
-                    sectionName["docname"] = partname
-                    startOfFile.append(sectionName)
-                    sectionName.append(title)
-                    compoundParent.append(startOfFile)
-                    for node in self.document.traverse(docutils.nodes.compound):
-                        if "toctree-wrapper" in node["classes"]:
-                            flag = checkNodeIsInPart(f, node)
-                            if flag:
-                                nodecopy = node
-                                replaceWithNode(node, HiddenCellNode, False)
-                                sectionName.append(nodecopy)
-                    self.document.append(compoundParent)
-
-            # append bib at the end
-            if len(bibNodes):
-                self.document.extend(bibNodes)
+        # move toctrees to the end of their parent section
+        for original_node in self.document.traverse(docutils.nodes.compound):
+            if "toctree-wrapper" in original_node["classes"]:
+                parent_node = find_parent(self.app.env, original_node, "section")
+                if parent_node:
+                    replace_node_cls(original_node, HiddenCellNode, False)
+                    parent_node.append(original_node)
+
+        # store bibliography nodes to append it at the end of doc
+        bib_nodes = []
+        for bib_node in self.document.traverse(thebibliography):
+            bib_nodes.append(bib_node)
+            replace_node_cls(bib_node, HiddenCellNode, False)
+
+        toc_path = Path(self.app.confdir or self.app.srcdir).joinpath("_toc.yml")
+        tocfile = yaml.safe_load(toc_path.read_text("utf8"))
+
+        for f in tocfile:
+            if "part" in f:
+                self.app.config["latex_toplevel_sectioning"] = "part"
+                partname = f["part"]
+                compound_parent = docutils.nodes.compound("")
+                compound_parent["classes"] = "toctree-wrapper"
+                start_of_file = addnodes.start_of_file("")
+                start_of_file["docname"] = partname
+                title = docutils.nodes.title(text=partname)
+                section_node = docutils.nodes.section("")
+                section_node["docname"] = partname
+                start_of_file.append(section_node)
+                section_node.append(title)
+                compound_parent.append(start_of_file)
+                for node in self.document.traverse(docutils.nodes.compound):
+                    if "toctree-wrapper" in node["classes"]:
+                        flag = check_node_in_part(f, node, self.app)
+                        if flag:
+                            original_node = node
+                            replace_node_cls(node, HiddenCellNode, False)
+                            section_node.append(original_node)
+                self.document.append(compound_parent)
+
+        # append bib at the end
+        if bib_nodes:
+            self.document.extend(bib_nodes)

setup.cfg

@@ -60,3 +60,5 @@ strict_equality = True
 [flake8]
 max-line-length = 100
 extend-ignore = E203
+per-file-ignores =
+    jupyterbook_latex/nodes.py:N802

tests/roots/test-partsToc/intro.md

@@ -9,3 +9,7 @@ Etiam interdum tempor augue at volutpat. Nulla sit amet volutpat elit. In vehicu
 ### This is H3
 
 Sed elit mi, vestibulum sit amet porttitor in, elementum a sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Donec gravida faucibus enim, at venenatis ante. Curabitur viverra neque vel nulla condimentum ultrices quis a ipsum.
+
+#### This is H4
+
+Ipso facto