Change [TOC] processing in MD browser to match gitiles

tools/md_browser/md_browser.py

Index: tools/md_browser/md_browser.py
diff --git a/tools/md_browser/md_browser.py b/tools/md_browser/md_browser.py
index 43dc5bdc921deae8911c6f72d1152d417d07bc8e..322c145ab979275f79a9d01a8b1f284384f5150f 100755
--- a/tools/md_browser/md_browser.py
+++ b/tools/md_browser/md_browser.py
@@ -17,11 +17,12 @@ import sys
 import threading
 import time
 import webbrowser
+from xml.etree import ElementTree
 
 
 THIS_DIR = os.path.realpath(os.path.dirname(__file__))
 SRC_DIR = os.path.dirname(os.path.dirname(THIS_DIR))
-sys.path.append(os.path.join(SRC_DIR, 'third_party', 'Python-Markdown'))
+sys.path.insert(0, os.path.join(SRC_DIR, 'third_party', 'Python-Markdown'))
 import markdown
 
 
@@ -155,10 +156,19 @@ class Handler(SimpleHTTPServer.SimpleHTTPRequestHandler):
     }
 
     contents = self._Read(path[1:])
-    md_fragment = markdown.markdown(contents,
-                                    extensions=extensions,
-                                    extension_configs=extension_configs,
-                                    output_format='html4').encode('utf-8')
+
+    md = markdown.Markdown(extensions=extensions,
+                           extension_configs=extension_configs,
+                           output_format='html4')
+
+    has_a_single_h1 = (len([line for line in contents.splitlines()
+                            if (line.startswith('#') and
+                                not line.startswith('##'))]) == 1)
+
+    md.treeprocessors['adjust_toc'] = _AdjustTOC(has_a_single_h1)
+
+    md_fragment = md.convert(contents).encode('utf-8')
+
     try:
       self._WriteHeader('text/html')
       self._WriteTemplate('header.html')
@@ -199,5 +209,70 @@ class Handler(SimpleHTTPServer.SimpleHTTPRequestHandler):
     self.wfile.write(contents.encode('utf-8'))
 
 
+class _AdjustTOC(markdown.treeprocessors.Treeprocessor):
+  def __init__(self, has_a_single_h1):
+    super(_AdjustTOC, self).__init__()
+    self.has_a_single_h1 = has_a_single_h1
+
+  def run(self, tree):
+    # The markdown.extensions.toc extension will replace
+    #
+    #     # H1
+    #
+    #     [TOC]
+    #
+    #     ## first H2
+    #
+    #     ## second H2
+    #
+    # with the following HTML:
+    #
+    #     <div class='toc'>
+    #       <ul><li><a>H1</a>
+    #               <ul><li>first H2
+    #                   <li>second H2</li></ul></li><ul></div>
+    #
+    # We want to replace the first ul with the second one, and also patch
+    # in a new H2 header for the TOC itself, ending up with:
+    #
+    #     <div class='toc'>
+    #       <h2>Contents</h2>
+    #       <div class='toc-aux'>
+    #       <ul><li>first H2
+    #           <li>second H2</li></ul></div></div>
+    #
+    # If the document has more than one H1, or no H1s at all, we leave
+    # the tree alone but still wrap it in a toc-aux div and add the
+    # Contents header.
+
+    toc_nodes = tree.findall(".//*[@class='toc']")
+    if not toc_nodes:
+      return
+
+    for toc_node in toc_nodes:
+      if self.has_a_single_h1:
+        toc_ul = toc_node[0]
+        toc_ul_li = toc_ul[0]
+        toc_ul_li_ul = toc_ul_li[1]
+        children = toc_ul_li_ul
+        del toc_node[0]
+      else:
+        children = list(toc_node)

[jsbell, 2016/11/08 17:29:12]

I'm unfamiliar with ElementTree and the exact structure in this case, but
assuming the structure is the same (<div><ul><li>...) it looks like since
toc_node is the <div>, children here would end up being just the <ul>, whereas
we want it to be the list of <li>s?

[Dirk Pranke, 2016/11/08 21:55:07]

Ah, good catch.

+        while len(toc_node):
+          del toc_node[0]
+
+      contents = ElementTree.SubElement(toc_node, 'h2')
+      contents.text = 'Contents'
+      contents.tail = '\n'
+      toc_aux = ElementTree.SubElement(toc_node, 'div',
+                                       attrib={'class': 'toc-aux'})
+      toc_aux.text = '\n'
+      toc_aux_ul = ElementTree.SubElement(toc_aux, 'ul')
+      toc_aux_ul.text = '\n'
+      toc_aux_ul.extend(children)
+      toc_aux_ul.tail = '\n'
+      toc_aux.tail = '\n'
+
+
 if __name__ == '__main__':
   sys.exit(main(sys.argv[1:]))