Getting closer to a table of contents part/decoration

toc_contents_docutils_part.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
import docutils
from docutils import io, languages, nodes
from docutils.core import (Publisher, publish_doctree,
                           publish_parts)
from docutils.nodes import Decorative, Element, decoration
from docutils.transforms import Transform, parts
from docutils.writers.html5_polyglot import HTMLTranslator, Writer


contents = """

=========
Hey world
=========

what's up
---------

ok



"""


def publish_parts_from_doctree(document, destination_path=None,
                               writer=None, writer_name='pseudoxml',
                               settings=None, settings_spec=None,
                               settings_overrides=None, config_section=None,
                               enable_exit_status=False):
    """
    Set up & run a `Publisher` to render from an existing document
    tree data structure, for programmatic use with string I/O.  Return
    the encoded string output.

    Note that document.settings is overridden; if you want to use the settings
    of the original `document`, pass settings=document.settings.

    Also, new document.transformer and document.reporter objects are
    generated.

    For encoded string output, be sure to set the 'output_encoding' setting to
    the desired encoding.  Set it to 'unicode' for unencoded Unicode string
    output.  Here's one way::

        publish_from_doctree(
            ..., settings_overrides={'output_encoding': 'unicode'})

    Parameters: `document` is a `docutils.nodes.document` object, an existing
    document tree.

    Other parameters: see `publish_programmatically`.
    """
    reader = docutils.readers.doctree.Reader(parser_name='null')
    pub = Publisher(reader, None, writer,
                    source=io.DocTreeInput(document),
                    destination_class=io.StringOutput, settings=settings)
    if not writer and writer_name:
        pub.set_writer(writer_name)
    pub.process_programmatic_settings(
        settings_spec, settings_overrides, config_section)
    pub.set_destination(None, destination_path)
    pub.publish(enable_exit_status=enable_exit_status)
    return pub.writer.parts


# parts for the contents
doc = publish_doctree(
    source=contents,
)


class toc(Decorative, Element):
    pass


class toc_decoration(decoration):

    def get_toc(self):
        print('get_toc decoration')
        if not len(self.children) or not isinstance(self.children[-1], toc):
            self.append(toc())
        return self.children[-1]


def get_decoration(self):
    if not self.decoration:
        self.decoration = toc_decoration()
        index = self.first_child_not_matching_class(nodes.Titular)
        if index is None:
            self.append(self.decoration)
        else:
            self.insert(index, self.decoration)
    return self.decoration


class TestTransform(Transform):

    default_priority = 940

    def apply(self):
        print("HI")
        # monkeypatch get_decoration in nodes to get decorator with toc
        self.document.get_decoration = get_decoration
        print('apply decoration')
        toc_nodes = self.generate_toc()
        print("toc_nodes: %s" % toc_nodes)
        if toc_nodes:
            decoration = self.document.get_decoration(self.document)
            toc = decoration.get_toc()
            toc.extend(toc_nodes)

    def generate_toc(self):
        language = languages.get_language(self.document.settings.language_code,
                                          self.document.reporter)
        name = language.labels['contents']
        title = nodes.title('', name)
        topic = nodes.topic('', title, classes=['contents'])
        name = nodes.fully_normalize_name(name)
        if not self.document.has_name(name):
            topic['names'].append(name)
        self.document.note_implicit_target(topic)
        pending = nodes.pending(parts.Contents)
        topic += pending
        self.document.note_pending(pending)

        return[topic]


class NewHTMLTranslator(HTMLTranslator):

    def __init__(self, document):
        HTMLTranslator.__init__(self, document)

        def node_pass(self, node):
            pass

        def ignore_node(self, node):
            raise nodes.SkipNode()
        HTMLTranslator.visit_toc_decoration = node_pass
        HTMLTranslator.depart_toc_decoration = node_pass
        HTMLTranslator.visit_toc = node_pass
        HTMLTranslator.depart_toc = node_pass
        HTMLTranslator.depart_pending = node_pass
        HTMLTranslator.visit_pending = node_pass
        self.toc = []

    def visit_pending(self, node):
        raise nodes.SkipNode()

    def depart_pending(self, node):
        raise nodes.SkipNode()

    def visit_toc(self, node):
        print("HEY VISIT TOC")
        self.context.append(len(self.body))

    def depart_toc(self, node):
        start = self.context.pop()
        toc = [self.starttag(node, 'div', CLASS='toc'),
               '<hr class="toc" />\n']
        toc.extend(self.body[start:])
        toc.append('\n</div>\n')
        self.toc.extend(toc)
        self.body_suffix[:0] = toc
        del self.body[start:]
        print("HEY DEPART TOC")

    def visit_toc_decoration(self, node):
        print("HEY VISIT TOC DECORATION")

    def depart_toc_decoration(self, node):
        print("HEY DEPART TOC DECORATION")


class NewWriter(Writer):

    visitor_attributes = Writer.visitor_attributes + ('toc', )

    def __init__(self):
        Writer.__init__(self)
        # I'd like to put this into the class attribute, but I think
        # somewhere up the Writer/Translator hierarchy are "old" python
        # classes. (e.g. Python =< 2.1 classes)
        self.translator_class = NewHTMLTranslator

    def get_transforms(self):
        return Writer.get_transforms(self) + [TestTransform]


w = NewWriter()
print(publish_parts_from_doctree(doc, writer=w)['toc'])
print(publish_parts(contents, writer=w)['toc'])