agoose77/sphinx-toctree-signature.py

## sphinx-toctree-signature.py
from sphinx.transforms import SphinxTransform
import sphinx.environment.collectors.toctree as toctree_collector
from sphinx import addnodes
from docutils import nodes


from typing import Any, Dict, List, Set, Tuple, Type, TypeVar, cast

from docutils import nodes
from docutils.nodes import Element, Node

from sphinx import addnodes
from sphinx.application import Sphinx
from sphinx.environment import BuildEnvironment
from sphinx.environment.adapters.toctree import TocTree
from sphinx.environment.collectors import EnvironmentCollector
from sphinx.locale import __
from sphinx.transforms import SphinxContentsFilter
from sphinx.util import logging, url_re

N = TypeVar("N")

logger = logging.getLogger(__name__)


class SignatureContentsFilter(SphinxContentsFilter):
    # This logic is obtuse, but what's happening here is that we set up the
    # .parent array to hold a _container_ around a text node.
    # The `SkipNode` exception: https://github.com/docutils/docutils/blob/ae4d18314a821e61a24dc0e4f29a691b7c3b656e/docutils/docutils/nodes.py#L2159-L2164
    # ensures that the `default_depart` method: https://github.com/docutils/docutils/blob/ae4d18314a821e61a24dc0e4f29a691b7c3b656e/docutils/docutils/nodes.py#L2127-L2130
    # isn't called.
    # It also ensures the container's children aren't visited - we want to stop here
    # `fullname` is the qualified name of the documented object.
    def visit_desc_signature(self, node):
        # Replace this element with a simple element wrapping text
        self.parent.append(nodes.Element("", nodes.Text(node.attributes["fullname"])))
        raise nodes.SkipNode


# Most of this is copied from Sphinx
class BetterTocTreeCollector(toctree_collector.TocTreeCollector):
    def process_doc(self, app: Sphinx, doctree: nodes.document) -> None:
        """Build a TOC from the doctree and store it in the inventory."""
        docname = app.env.docname
        numentries = [0]  # nonlocal again...

        # This is changed to a generator, and the class condition removed
        def traverse_in_section(node: Element) -> List[N]:
            """Like traverse(), but stay within the same section."""
            yield node
            for child in node.children:
                if isinstance(child, nodes.section):
                    continue
                elif isinstance(child, nodes.Element):
                    yield from traverse_in_section(child)

        def build_toc(node: Element, depth: int = 1) -> nodes.bullet_list:
            # The logic here is a bit confusing.
            # It looks like section nodes are expected to be top-level within a section.
            entries: List[Element] = []

            for sectionnode in node:
                # find all toctree nodes in this section and add them
                # to the toc (just copying the toctree node which is then
                # resolved in self.get_and_resolve_doctree)
                if isinstance(sectionnode, nodes.section):
                    title = sectionnode[0]
                    # copy the contents of the section title, but without references
                    # and unnecessary stuff
                    visitor = SphinxContentsFilter(doctree)
                    title.walkabout(visitor)
                    nodetext = visitor.get_entry_text()
                    # if nodetext and nodetext[0] == "ak.ArrayBuilder":
                    # print(node)
                    # break
                    if not numentries[0]:
                        # for the very first toc entry, don't add an anchor
                        # as it is the file's title anyway
                        anchorname = ""
                    else:
                        anchorname = "#" + sectionnode["ids"][0]
                    numentries[0] += 1
                    # make these nodes:
                    # list_item -> compact_paragraph -> reference
                    reference = nodes.reference(
                        "",
                        "",
                        internal=True,
                        refuri=docname,
                        anchorname=anchorname,
                        *nodetext
                    )
                    para = addnodes.compact_paragraph("", "", reference)
                    item: Element = nodes.list_item("", para)
                    sub_item = build_toc(sectionnode, depth + 1)
                    if sub_item:
                        item += sub_item
                    entries.append(item)
                elif isinstance(sectionnode, addnodes.only):
                    onlynode = addnodes.only(expr=sectionnode["expr"])
                    blist = build_toc(sectionnode, depth)
                    if blist:
                        onlynode += blist.children
                        entries.append(onlynode)
                # Otherwise, for a generic element we allow recursion into the section
                elif isinstance(sectionnode, nodes.Element):
                    for node in traverse_in_section(sectionnode):
                        if isinstance(node, addnodes.toctree):
                            item = node.copy()
                            entries.append(item)
                            # important: do the inventory stuff
                            TocTree(app.env).note(docname, node)
                        # For signatures within some section, we add them to the ToC
                        elif isinstance(node, addnodes.desc_signature):
                            title = node
                            # Copy the signature, but without the detail e.g. parens
                            visitor = SignatureContentsFilter(doctree)
                            title.walkabout(visitor)
                            nodetext = visitor.get_entry_text()

                            if not numentries[0]:
                                # for the very first toc entry, don't add an anchor
                                # as it is the file's title anyway
                                anchorname = ""
                            else:
                                anchorname = "#" + node["ids"][0]
                            numentries[0] += 1
                            # make these nodes:
                            # list_item -> compact_paragraph -> reference
                            reference = nodes.reference(
                                "",
                                "",
                                internal=True,
                                refuri=docname,
                                anchorname=anchorname,
                                *nodetext
                            )
                            para = addnodes.compact_paragraph("", "", reference)
                            item: Element = nodes.list_item("", para)
                            entries.append(item)
            if entries:
                return nodes.bullet_list("", *entries)
            return None

        toc = build_toc(doctree)
        assert docname in app.env.tocs
        if toc:
            app.env.tocs[docname] = toc
        else:
            app.env.tocs[docname] = nodes.bullet_list("")
        app.env.toc_num_entries[docname] = numentries[0]


def setup(app):
    app.add_env_collector(BetterTocTreeCollector)
	from sphinx.transforms import SphinxTransform
	import sphinx.environment.collectors.toctree as toctree_collector
	from sphinx import addnodes
	from docutils import nodes


	from typing import Any, Dict, List, Set, Tuple, Type, TypeVar, cast

	from docutils import nodes
	from docutils.nodes import Element, Node

	from sphinx import addnodes
	from sphinx.application import Sphinx
	from sphinx.environment import BuildEnvironment
	from sphinx.environment.adapters.toctree import TocTree
	from sphinx.environment.collectors import EnvironmentCollector
	from sphinx.locale import __
	from sphinx.transforms import SphinxContentsFilter
	from sphinx.util import logging, url_re

	N = TypeVar("N")

	logger = logging.getLogger(__name__)


	class SignatureContentsFilter(SphinxContentsFilter):
	# This logic is obtuse, but what's happening here is that we set up the
	# .parent array to hold a _container_ around a text node.
	# The `SkipNode` exception: https://github.com/docutils/docutils/blob/ae4d18314a821e61a24dc0e4f29a691b7c3b656e/docutils/docutils/nodes.py#L2159-L2164
	# ensures that the `default_depart` method: https://github.com/docutils/docutils/blob/ae4d18314a821e61a24dc0e4f29a691b7c3b656e/docutils/docutils/nodes.py#L2127-L2130
	# isn't called.
	# It also ensures the container's children aren't visited - we want to stop here
	# `fullname` is the qualified name of the documented object.
	def visit_desc_signature(self, node):
	# Replace this element with a simple element wrapping text
	self.parent.append(nodes.Element("", nodes.Text(node.attributes["fullname"])))
	raise nodes.SkipNode


	# Most of this is copied from Sphinx
	class BetterTocTreeCollector(toctree_collector.TocTreeCollector):
	def process_doc(self, app: Sphinx, doctree: nodes.document) -> None:
	"""Build a TOC from the doctree and store it in the inventory."""
	docname = app.env.docname
	numentries = [0] # nonlocal again...

	# This is changed to a generator, and the class condition removed
	def traverse_in_section(node: Element) -> List[N]:
	"""Like traverse(), but stay within the same section."""
	yield node
	for child in node.children:
	if isinstance(child, nodes.section):
	continue
	elif isinstance(child, nodes.Element):
	yield from traverse_in_section(child)

	def build_toc(node: Element, depth: int = 1) -> nodes.bullet_list:
	# The logic here is a bit confusing.
	# It looks like section nodes are expected to be top-level within a section.
	entries: List[Element] = []

	for sectionnode in node:
	# find all toctree nodes in this section and add them
	# to the toc (just copying the toctree node which is then
	# resolved in self.get_and_resolve_doctree)
	if isinstance(sectionnode, nodes.section):
	title = sectionnode[0]
	# copy the contents of the section title, but without references
	# and unnecessary stuff
	visitor = SphinxContentsFilter(doctree)
	title.walkabout(visitor)
	nodetext = visitor.get_entry_text()
	# if nodetext and nodetext[0] == "ak.ArrayBuilder":
	# print(node)
	# break
	if not numentries[0]:
	# for the very first toc entry, don't add an anchor
	# as it is the file's title anyway
	anchorname = ""
	else:
	anchorname = "#" + sectionnode["ids"][0]
	numentries[0] += 1
	# make these nodes:
	# list_item -> compact_paragraph -> reference
	reference = nodes.reference(
	"",
	"",
	internal=True,
	refuri=docname,
	anchorname=anchorname,
	*nodetext
	)
	para = addnodes.compact_paragraph("", "", reference)
	item: Element = nodes.list_item("", para)
	sub_item = build_toc(sectionnode, depth + 1)
	if sub_item:
	item += sub_item
	entries.append(item)
	elif isinstance(sectionnode, addnodes.only):
	onlynode = addnodes.only(expr=sectionnode["expr"])
	blist = build_toc(sectionnode, depth)
	if blist:
	onlynode += blist.children
	entries.append(onlynode)
	# Otherwise, for a generic element we allow recursion into the section
	elif isinstance(sectionnode, nodes.Element):
	for node in traverse_in_section(sectionnode):
	if isinstance(node, addnodes.toctree):
	item = node.copy()
	entries.append(item)
	# important: do the inventory stuff
	TocTree(app.env).note(docname, node)
	# For signatures within some section, we add them to the ToC
	elif isinstance(node, addnodes.desc_signature):
	title = node
	# Copy the signature, but without the detail e.g. parens
	visitor = SignatureContentsFilter(doctree)
	title.walkabout(visitor)
	nodetext = visitor.get_entry_text()

	if not numentries[0]:
	# for the very first toc entry, don't add an anchor
	# as it is the file's title anyway
	anchorname = ""
	else:
	anchorname = "#" + node["ids"][0]
	numentries[0] += 1
	# make these nodes:
	# list_item -> compact_paragraph -> reference
	reference = nodes.reference(
	"",
	"",
	internal=True,
	refuri=docname,
	anchorname=anchorname,
	*nodetext
	)
	para = addnodes.compact_paragraph("", "", reference)
	item: Element = nodes.list_item("", para)
	entries.append(item)
	if entries:
	return nodes.bullet_list("", *entries)
	return None

	toc = build_toc(doctree)
	assert docname in app.env.tocs
	if toc:
	app.env.tocs[docname] = toc
	else:
	app.env.tocs[docname] = nodes.bullet_list("")
	app.env.toc_num_entries[docname] = numentries[0]


	def setup(app):
	app.add_env_collector(BetterTocTreeCollector)