Change [TOC] processing in MD browser to match gitiles

tools/md_browser/md_browser.py

 #!/usr/bin/env python
 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Simple Markdown browser for a Git checkout."""
 from __future__ import print_function
 
 import SimpleHTTPServer
 import SocketServer
 import argparse
 import codecs
 import os
 import re
 import socket
 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
 
 
 def main(argv):
   parser = argparse.ArgumentParser(prog='md_browser')
   parser.add_argument('-p', '--port', type=int, default=8080,
                       help='port to run on (default = %(default)s)')
   parser.add_argument('-d', '--directory', type=str, default=SRC_DIR)
   parser.add_argument('file', nargs='?',
                       help='open file in browser')
@@ skipping 113 matching lines @@
         'markdown.extensions.toc',
         'gitiles_ext_blocks',
     ]
     extension_configs = {
         'markdown.extensions.toc': {
             'slugify': _gitiles_slugify
         },
     }
 
     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')
       self.wfile.write(md_fragment)
       self._WriteTemplate('footer.html')
     except:
       raise
 
   def _DoCSS(self, template):
     self._WriteHeader('text/css')
@@ skipping 20 matching lines @@
     self.send_response(200)
     self.send_header('Content-Type', content_type)
     self.end_headers()
 
   def _WriteTemplate(self, template):
     contents = self._Read(os.path.join('tools', 'md_browser', template),
                           relative_to=SRC_DIR)
     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:]))