From 0242a96366b1681a86563d63776e320ff6f76a9e Mon Sep 17 00:00:00 2001
From: Charles de Beauchesne <charles.de.beauchesne@gmail.com>
Date: Thu, 8 Mar 2018 14:56:09 +0100
Subject: Better check of allowed TOC location #639 (#641)

---
 markdown/extensions/toc.py | 22 ++++++++++++----------
 tests/test_extensions.py   | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/markdown/extensions/toc.py b/markdown/extensions/toc.py
index 21f98ba..0e9b2b9 100644
--- a/markdown/extensions/toc.py
+++ b/markdown/extensions/toc.py
@@ -137,11 +137,17 @@ class TocTreeprocessor(Treeprocessor):
 
         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 @@ class TocTreeprocessor(Treeprocessor):
             # 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
index 5a04e64..aee9bac 100644
--- a/tests/test_extensions.py
+++ b/tests/test_extensions.py
@@ -845,6 +845,41 @@ class TestTOC(TestCaseWithAssertStartsWith):
         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):
-- 
cgit v1.2.3