Better check of allowed TOC location Python-Markdown#639 (Python-Mark…

…down#641)
x0x1x0x · Mar 8, 2018 · 0242a96 · 0242a96
1 parent cab4b69
commit 0242a96
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 10 deletions.
diff --git a/markdown/extensions/toc.py b/markdown/extensions/toc.py
@@ -137,11 +137,17 @@ def __init__(self, md, config):
 
         self.header_rgx = re.compile("[Hh][123456]")
 
-    def iterparent(self, root):
-        ''' Iterator wrapper to get parent and child all at once. '''
-        for parent in root.iter():
-            for child in parent:
-                yield parent, child
+    def iterparent(self, node):
+        ''' Iterator wrapper to get allowed parent and child all at once. '''
+
+        # We do not allow the marker inside a header as that
+        # would causes an enless loop of placing a new TOC
+        # inside previously generated TOC.
+        for child in node:
+            if not self.header_rgx.match(child.tag) and child.tag not in ['pre', 'code']:
+                yield node, child
+                for p, c in self.iterparent(child):
+                    yield p, c
 
     def replace_marker(self, root, elem):
         ''' Replace marker with elem. '''
@@ -153,11 +159,7 @@ def replace_marker(self, root, elem):
             # To keep the output from screwing up the
             # validation by putting a <div> inside of a <p>
             # we actually replace the <p> in its entirety.
-            # We do not allow the marker inside a header as that
-            # would causes an enless loop of placing a new TOC
-            # inside previously generated TOC.
-            if c.text and c.text.strip() == self.marker and \
-               not self.header_rgx.match(c.tag) and c.tag not in ['pre', 'code']:
+            if c.text and c.text.strip() == self.marker:
                 for i in range(len(p)):
                     if p[i] == c:
                         p[i] = elem

diff --git a/tests/test_extensions.py b/tests/test_extensions.py
@@ -845,6 +845,41 @@ def testUniqueFunc(self):
         self.assertEqual(unique('foo', ids), 'foo_1')
         self.assertEqual(ids, set(['foo', 'foo_1']))
 
+    def testTocInHeaders(self):
+
+        text = '[TOC]\n#[TOC]'
+        self.assertEqual(
+            self.md.convert(text),
+            '<div class="toc">\n'                       # noqa
+              '<ul>\n'                                  # noqa
+                '<li><a href="#toc">[TOC]</a></li>\n'   # noqa
+              '</ul>\n'                                 # noqa
+            '</div>\n'                                  # noqa
+            '<h1 id="toc">[TOC]</h1>'                   # noqa
+        )
+
+        text = '#[TOC]\n[TOC]'
+        self.assertEqual(
+            self.md.convert(text),
+            '<h1 id="toc">[TOC]</h1>\n'                 # noqa
+            '<div class="toc">\n'                       # noqa
+              '<ul>\n'                                  # noqa
+                '<li><a href="#toc">[TOC]</a></li>\n'   # noqa
+              '</ul>\n'                                 # noqa
+            '</div>'                                    # noqa
+        )
+
+        text = '[TOC]\n# *[TOC]*'
+        self.assertEqual(
+            self.md.convert(text),
+            '<div class="toc">\n'                       # noqa
+              '<ul>\n'                                  # noqa
+                '<li><a href="#toc">[TOC]</a></li>\n'   # noqa
+              '</ul>\n'                                 # noqa
+            '</div>\n'                                  # noqa
+            '<h1 id="toc"><em>[TOC]</em></h1>'          # noqa
+        )
+
 
 class TestSmarty(unittest.TestCase):
     def setUp(self):