Merge pull request #59 from jbms/jbms-minor-fixes

Speed up TOC generation and other minor fixes

sphinx_immaterial/__init__.py

@@ -1,12 +1,13 @@
 """Sphinx-Immaterial theme."""
 
 import os
-from typing import List, Type, Dict, Mapping
+from typing import cast, List, Type, Dict, Mapping, Optional
 
 import docutils.nodes
 from sphinx.application import Sphinx
 import sphinx.builders
 import sphinx.builders.html
+import sphinx.theming
 import sphinx.util.logging
 import sphinx.util.fileutil
 import sphinx.util.matching
@@ -58,6 +59,12 @@ def __init__(self, *args, **kwargs):
                 self.settings.table_style.split(",") + ["data"]
             )
 
+            # Ensure classes like `s` (used for string literals in code
+            # highlighting) aren't converted to `<s>` elements (strikethrough).
+            # Sphinx already overrides this, but for some reason due to
+            # `__init__` invocation order it gets overridden.
+            self.supported_inline_tags = set()
+
         def visit_section(self, node: docutils.nodes.section) -> None:
             # Sphinx normally writes sections with a section heading as:
             #
@@ -92,6 +99,11 @@ def _get_html_builder(base_builder: Type[sphinx.builders.html.StandaloneHTMLBuil
     """Returns a modified HTML translator."""
 
     class CustomHTMLBuilder(base_builder):
+
+        css_files: List[sphinx.builders.html.Stylesheet]
+        theme: sphinx.theming.Theme
+        templates: sphinx.jinja2glue.BuiltinTemplateLoader
+
         @property
         def default_translator_class(self):
             return _get_html_translator(super().default_translator_class)
@@ -127,7 +139,11 @@ def init_css_files(self):
                     "_static/basic.css",
                 ]
             )
-            self.css_files = [x for x in self.css_files if x.filename not in excluded]
+            self.css_files = [
+                x
+                for x in cast(List[sphinx.builders.html.Stylesheet], self.css_files)
+                if x.filename not in excluded
+            ]
 
         def gen_additional_pages(self):
             # Prevent the search.html page from being written since this theme provides
@@ -171,11 +187,13 @@ def onerror(filename: str, error: Exception) -> None:
                         os.path.join(self.outdir, "_static"),
                         excluded=excluded,
                         context=context,
-                        renderer=self.templates,
+                        renderer=cast(
+                            sphinx.util.template.BaseRenderer, self.templates
+                        ),
                         onerror=onerror,
                     )
 
-        def get_target_uri(self, docname: str, typ: str = None) -> str:
+        def get_target_uri(self, docname: str, typ: Optional[str] = None) -> str:
             """Strips ``index.html`` suffix from URIs for cleaner links."""
             orig_uri = super().get_target_uri(docname, typ)
             if self.app.config["html_use_directory_uris_for_index_pages"]:
@@ -189,7 +207,7 @@ def get_target_uri(self, docname: str, typ: str = None) -> str:
     return CustomHTMLBuilder
 
 
-def dict_merge(*dicts: List[Mapping]):
+def dict_merge(*dicts: Mapping):
     """Recursively merges the members of one or more dicts."""
     result = {}
     for d in dicts:

sphinx_immaterial/content_tabs.py

@@ -172,8 +172,12 @@ def depart_tab_set(self, node):
     self.body.append("</div>")
 
 
-def setup(app: Sphinx) -> None:
+def setup(app: Sphinx):
     app.add_directive("md-tab-set", MaterialTabSetDirective)
     app.add_directive("md-tab-item", MaterialTabItemDirective)
     app.add_node(content_tab_label, html=(visit_tab_label, depart_tab_label))
     app.add_node(content_tab_set, html=(visit_tab_set, depart_tab_set))
+    return {
+        "parallel_read_safe": True,
+        "parallel_write_safe": True,
+    }

sphinx_immaterial/details_patch.py

@@ -2,7 +2,9 @@
 from docutils import nodes
 
 try:
-    from sphinxcontrib.details.directive import DetailsDirective
+    from sphinxcontrib.details.directive import (  # pytype: disable=import-error
+        DetailsDirective,
+    )
 except ImportError:
     DetailsDirective = None
 

sphinx_immaterial/md_admonition.py

@@ -12,15 +12,20 @@ class NoTitleAdmonition(admonitions.BaseAdmonition):
 
     optional_arguments = 1
     node_class = nodes.admonition
+    default_title = ""
+
+    classes = ()
 
     def run(self):
         set_classes(self.options)
+        if self.classes:
+            self.options.setdefault("classes", list(self.classes))
         self.assert_has_content()
         text = "\n".join(self.content)
         admonition_node = self.node_class(text, **self.options)
         self.add_name(admonition_node)
         if self.node_class is nodes.admonition:
-            title_text = self.arguments[0] if self.arguments else ""
+            title_text = self.arguments[0] if self.arguments else self.default_title
             textnodes, messages = self.state.inline_text(title_text, self.lineno)
             title = nodes.title(title_text, "", *textnodes)
             title.source, title.line = self.state_machine.get_source_and_line(
@@ -29,7 +34,7 @@ def run(self):
             if title_text:
                 admonition_node += title
             admonition_node += messages
-            if not "classes" in self.options and title_text:
+            if not self.options.get("classes") and title_text:
                 admonition_node["classes"] += ["admonition" + nodes.make_id(title_text)]
         self.state.nested_parse(self.content, self.content_offset, admonition_node)
         return [admonition_node]

sphinx_immaterial/mermaid_diagrams.py

@@ -61,3 +61,7 @@ def setup(app: Sphinx):
         html=(visit_mermaid_node_html, depart_mermaid_node_html),
         latex=(visit_mermaid_node_latex, depart_mermaid_node_latex),
     )
+    return {
+        "parallel_read_safe": True,
+        "parallel_write_safe": True,
+    }

sphinx_immaterial/nav_adapt.py

@@ -1,15 +1,27 @@
 """Injects mkdocs-style `nav` and `page` objects into the HTML jinja2 context."""
 
+import collections
 import copy
 import os
 import re
-from typing import List, Union, NamedTuple, Optional, Tuple, Iterator, Dict
+from typing import (
+    List,
+    Union,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Iterator,
+    Dict,
+    Iterable,
+    Set,
+)
 import urllib.parse
 import docutils.nodes
 import markupsafe
 import sphinx.builders
 import sphinx.application
 import sphinx.environment.adapters.toctree
+import sphinx.util.osutil
 
 
 # env var is only defined in RTD hosted builds
@@ -68,7 +80,9 @@ class _TocVisitor(docutils.nodes.NodeVisitor):
     """NodeVisitor used by `_get_mkdocs_toc`."""
 
     def __init__(
-        self, document: docutils.nodes.document, builder: sphinx.builders.Builder
+        self,
+        document: docutils.nodes.document,
+        builder: sphinx.builders.html.StandaloneHTMLBuilder,
     ):
         super().__init__(document)
         self._prev_caption: Optional[docutils.nodes.Element] = None
@@ -168,7 +182,7 @@ def visit_list_item(self, node: docutils.nodes.list_item):
 
 
 def _get_mkdocs_toc(
-    toc_node: docutils.nodes.Node, builder: sphinx.builders.Builder
+    toc_node: docutils.nodes.Node, builder: sphinx.builders.html.StandaloneHTMLBuilder
 ) -> List[MkdocsNavEntry]:
     """Converts a docutils toc node into a mkdocs-format JSON toc."""
     visitor = _TocVisitor(sphinx.util.docutils.new_document(""), builder)
@@ -258,6 +272,8 @@ class ObjectIconInfo(NamedTuple):
     ("cpp", "var"): ObjectIconInfo(icon_class="alias", icon_text="V"),
     ("cpp", "type"): ObjectIconInfo(icon_class="alias", icon_text="T"),
     ("cpp", "namespace"): ObjectIconInfo(icon_class="alias", icon_text="N"),
+    ("cpp", "templateParam"): ObjectIconInfo(icon_class="data", icon_text="P"),
+    ("cpp", "functionParam"): ObjectIconInfo(icon_class="sub-data", icon_text="p"),
 }
 
 
@@ -289,10 +305,10 @@ def _get_domain_anchor_map(
     app: sphinx.application.Sphinx,
 ) -> Dict[Tuple[str, str], DomainAnchorEntry]:
     key = "sphinx_immaterial_domain_anchor_map"
-    m = app.env.temp_data.get(key)
+    m = getattr(app.env, key, None)
     if m is None:
         m = _make_domain_anchor_map(app.env)
-        app.env.temp_data[key] = m
+        setattr(app.env, key, m)
     return m
 
 
@@ -359,20 +375,141 @@ def _collapse_children_not_on_same_page(entry: MkdocsNavEntry) -> MkdocsNavEntry
     return entry
 
 
+TocEntryKey = Tuple[int, ...]
+
+
+def _build_toc_index(toc: List[MkdocsNavEntry]) -> Dict[str, List[TocEntryKey]]:
+    """Builds a map from URL to list of toc entries.
+
+    This is used by `_get_global_toc` to efficiently prune the cached TOC for a
+    given page.
+    """
+    url_map: Dict[str, List[TocEntryKey]] = collections.defaultdict(list)
+
+    def _traverse(entries: List[MkdocsNavEntry], parent_key: TocEntryKey):
+        for i, entry in enumerate(entries):
+            child_key = parent_key + (i,)
+            url = entry.url
+            if url is None:
+                continue
+            url = _strip_fragment(url)
+            url_map[url].append(child_key)
+            _traverse(entry.children, child_key)
+
+    _traverse(toc, ())
+    return url_map
+
+
+class CachedTocInfo:
+    """Cached representation of the global TOC.
+
+    This is generated once (per writer process) and re-used for all pages.
+
+    Obtaining the global TOC via `TocTree.get_toctree_for` is expensive because
+    we first have to obtain a complete TOC of all pages, and then prune it for
+    the current page.  The overall cost to generate the TOCs for all pages
+    therefore ends up being quadratic in the number of pages, and in practice as
+    the number of pages reaches several hundred, a significant fraction of the
+    total documentation generation time is due to the TOC.
+
+    By the caching the TOC and a `url_map` that allows to efficiently prune the
+    TOC for a given page, the cost of generating the TOC for each page is much
+    lower.
+    """
+
+    def __init__(self, app: sphinx.application.Sphinx):
+        # Sphinx always generates a TOC relative to a particular page, and
+        # converts all page references to relative URLs.  Use an empty string as
+        # the page name to ensure the relative URLs that Sphinx generates are
+        # relative to the base URL.  When generating the per-page TOCs from this
+        # cached data structure the URLs will be converted to be relative to the
+        # current page.
+        fake_pagename = ""
+        global_toc_node = sphinx.environment.adapters.toctree.TocTree(
+            app.env
+        ).get_toctree_for(
+            fake_pagename,
+            app.builder,
+            collapse=False,
+            maxdepth=-1,
+            titles_only=False,
+        )
+        global_toc = _get_mkdocs_toc(global_toc_node, app.builder)
+        _add_domain_info_to_toc(app, global_toc, fake_pagename)
+        self.entries = global_toc
+        self.url_map = _build_toc_index(global_toc)
+
+
+def _get_cached_globaltoc_info(app: sphinx.application.Sphinx) -> CachedTocInfo:
+    """Obtains the cached global TOC, generating it if necessary."""
+    key = "sphinx_immaterial_global_toc_cache"
+    data = getattr(app.env, key, None)
+    if data is not None:
+        return data
+    data = CachedTocInfo(app)
+    setattr(app.env, key, data)
+    return data
+
+
+def _get_ancestor_keys(keys: Iterable[TocEntryKey]) -> Set[TocEntryKey]:
+    ancestors = set()
+    for key in keys:
+        while key not in ancestors:
+            ancestors.add(key)
+            key = key[:-1]
+    return ancestors
+
+
+def _get_global_toc(app: sphinx.application.Sphinx, pagename: str, collapse: bool):
+    """Obtains the global TOC for a given page."""
+    cached_data = _get_cached_globaltoc_info(app)
+    url = app.builder.get_target_uri(pagename)
+    keys = set(cached_data.url_map[url])
+    ancestors = _get_ancestor_keys(keys)
+
+    fake_pagename = ""
+
+    fake_page_url = app.builder.get_target_uri(fake_pagename)
+
+    real_page_url = app.builder.get_target_uri(pagename)
+
+    def _make_toc_for_page(key: TocEntryKey, children: List[MkdocsNavEntry]):
+        children = list(children)
+        for i, child in enumerate(children):
+            child_key = key + (i,)
+            child = children[i] = copy.copy(child)
+            if child.url is not None:
+                root_relative_url = urllib.parse.urljoin(fake_page_url, child.url)
+                uri = urllib.parse.urlparse(root_relative_url)
+                if not uri.netloc:
+                    child.url = sphinx.util.osutil.relative_uri(
+                        real_page_url, root_relative_url
+                    )
+                    if uri.fragment:
+                        child.url += f"#{uri.fragment}"
+            in_ancestors = child_key in ancestors
+            if in_ancestors:
+                child.active = True
+                if child_key in keys:
+                    child.current = True
+            if in_ancestors or child.caption_only:
+                child.children = _make_toc_for_page(child_key, child.children)
+            else:
+                child.children = []
+        return children
+
+    return _make_toc_for_page((), cached_data.entries)
+
+
 def _get_mkdocs_tocs(
     app: sphinx.application.Sphinx, pagename: str, duplicate_local_toc: bool
 ) -> Tuple[List[MkdocsNavEntry], List[MkdocsNavEntry]]:
     theme_options = app.config["html_theme_options"]
-    global_toc_node = sphinx.environment.adapters.toctree.TocTree(
-        app.env
-    ).get_toctree_for(
-        pagename,
-        app.builder,
+    global_toc = _get_global_toc(
+        app=app,
+        pagename=pagename,
         collapse=theme_options.get("globaltoc_collapse", False),
-        maxdepth=-1,
-        titles_only=False,
     )
-    global_toc = _get_mkdocs_toc(global_toc_node, app.builder)
     local_toc = []
     if pagename != app.env.config.master_doc:
         # Extract entry from `global_toc` corresponding to the current page.
@@ -392,9 +529,7 @@ def _get_mkdocs_tocs(
             app.builder,
         )
         local_toc = _get_mkdocs_toc(local_toc_node, app.builder)
-
-    _add_domain_info_to_toc(app, global_toc, pagename)
-    _add_domain_info_to_toc(app, local_toc, pagename)
+        _add_domain_info_to_toc(app, local_toc, pagename)
 
     if len(local_toc) == 1 and len(local_toc[0].children) == 0:
         local_toc = []
@@ -409,7 +544,7 @@ def _html_page_context(
     context: dict,
     doctree: docutils.nodes.Node,
 ) -> None:
-    theme_options = app.config["html_theme_options"]  # type: dict
+    theme_options: dict = app.config["html_theme_options"]
     page_title = markupsafe.Markup.escape(
         markupsafe.Markup(context.get("title")).striptags()
     )
@@ -478,8 +613,8 @@ def _html_page_context(
             ),
             "url": context["prev"]["link"],
         }
-    repo_url = theme_options.get("repo_url")  # type: str
-    edit_uri = theme_options.get("edit_uri")  # type: str
+    repo_url: Optional[str] = theme_options.get("repo_url")
+    edit_uri: Optional[str] = theme_options.get("edit_uri")
     if repo_url and edit_uri and not READTHEDOCS:
         page["edit_url"] = "/".join(
             [