From 3bca78cccf4ecc30b84b4043c1dab472aad749b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jes=C3=BAs=20Fern=C3=A1ndez?= <fernandez.cuesta@gmail.com>
Date: Tue, 1 Sep 2015 17:29:37 +0200
Subject: [PATCH] Add toc_depth parameter to toc extension

---
 docs/extensions/toc.txt    | 10 ++++++-
 markdown/extensions/toc.py | 11 +++++---
 tests/test_extensions.py   | 54 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/docs/extensions/toc.txt b/docs/extensions/toc.txt
index 2cd1fa15f..742759028 100644
--- a/docs/extensions/toc.txt
+++ b/docs/extensions/toc.txt
@@ -124,4 +124,12 @@ The following options are provided to configure the output:
     The callable must return a string appropriate for use in HTML `id` attributes.
 
 * **`separator`**:
-    Word separator. Character which replaces white space in id. Defaults to "`-`".
\ No newline at end of file
+    Word separator. Character which replaces white space in id. Defaults to "`-`".
+
+* **`toc_depth`**
+    Define up to which section level "n" (`<h1>` to `<hn>`, where `1 <= n <= 6`)
+    to include in the Table of Contents. Defaults to `6`.
+
+    When used with conjunction with `baselevel` this parameter will limit the
+    resulting (adjusted) heading. That is, if both `toc_depth` and `baselevel`
+    are 3, then only the highest level will be present in the table.
diff --git a/markdown/extensions/toc.py b/markdown/extensions/toc.py
index b3cf898f4..294a3cb33 100644
--- a/markdown/extensions/toc.py
+++ b/markdown/extensions/toc.py
@@ -136,8 +136,8 @@ def __init__(self, md, config):
         self.use_permalinks = parseBoolValue(config["permalink"], False)
         if self.use_permalinks is None:
             self.use_permalinks = config["permalink"]
-
         self.header_rgx = re.compile("[Hh][123456]")
+        self.toc_depth = config["toc_depth"]
 
     def iterparent(self, root):
         ''' Iterator wrapper to get parent and child all at once. '''
@@ -156,7 +156,7 @@ def replace_marker(self, root, elem):
             # 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
+            # would causes an endless 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']:
@@ -233,6 +233,8 @@ def run(self, doc):
         for el in doc.iter():
             if isinstance(el.tag, string_type) and self.header_rgx.match(el.tag):
                 self.set_level(el)
+                if int(el.tag[-1]) > int(self.toc_depth):
+                    continue
                 text = ''.join(el.itertext()).strip()
 
                 # Do not override pre-existing ids
@@ -284,7 +286,10 @@ def __init__(self, *args, **kwargs):
             "slugify": [slugify,
                         "Function to generate anchors based on header text - "
                         "Defaults to the headerid ext's slugify function."],
-            'separator': ['-', 'Word separator. Defaults to "-".']
+            'separator': ['-', 'Word separator. Defaults to "-".'],
+            "toc_depth": [6,
+                          "Define up to which section level n (<h1>..<hn>) to "
+                          "include in the TOC"]
         }
 
         super(TocExtension, self).__init__(*args, **kwargs)
diff --git a/tests/test_extensions.py b/tests/test_extensions.py
index 38f0be3c9..bb65d77d0 100644
--- a/tests/test_extensions.py
+++ b/tests/test_extensions.py
@@ -825,6 +825,60 @@ def testUniqueFunc(self):
         self.assertEqual(unique('foo', ids), 'foo_1')
         self.assertEqual(ids, set(['foo', 'foo_1']))
 
+    def testMaxLevel(self):
+        """ Test toc_depth setting """
+        md = markdown.Markdown(
+            extensions=[markdown.extensions.toc.TocExtension(toc_depth=2)]
+        )
+        text = '# Header 1\n\n## Header 2\n\n###Header 3 not in TOC'
+        self.assertEqual(
+            md.convert(text),
+            '<h1 id="header-1">Header 1</h1>\n'
+            '<h2 id="header-2">Header 2</h2>\n'
+            '<h3>Header 3 not in TOC</h3>'
+        )
+        self.assertEqual(
+            md.toc,
+            '<div class="toc">\n'
+              '<ul>\n'                                             # noqa
+                '<li><a href="#header-1">Header 1</a>'             # noqa
+                  '<ul>\n'                                         # noqa
+                    '<li><a href="#header-2">Header 2</a></li>\n'  # noqa
+                  '</ul>\n'                                        # noqa
+                '</li>\n'                                          # noqa
+              '</ul>\n'                                            # noqa
+            '</div>\n'
+        )
+
+        self.assertNotIn("Header 3", md.toc)
+
+    def testMaxLevelwithBaseLevel(self):
+        """ Test toc_depth setting together with baselevel """
+        md = markdown.Markdown(
+            extensions=[markdown.extensions.toc.TocExtension(toc_depth=3,
+                                                             baselevel=2)]
+        )
+        text = '# Some Header\n\n## Next Level\n\n### Too High'
+        self.assertEqual(
+            md.convert(text),
+            '<h2 id="some-header">Some Header</h2>\n'
+            '<h3 id="next-level">Next Level</h3>\n'
+            '<h4>Too High</h4>'
+        )
+        self.assertEqual(
+            md.toc,
+            '<div class="toc">\n'
+              '<ul>\n'                                                 # noqa
+                '<li><a href="#some-header">Some Header</a>'           # noqa
+                  '<ul>\n'                                             # noqa
+                    '<li><a href="#next-level">Next Level</a></li>\n'  # noqa
+                  '</ul>\n'                                            # noqa
+                '</li>\n'                                              # noqa
+              '</ul>\n'                                                # noqa
+            '</div>\n'
+        )
+        self.assertNotIn("Too High", md.toc)
+
 
 class TestSmarty(unittest.TestCase):
     def setUp(self):