Implement #43 Allow table of contents to generate from arbitrary depth of headers

mdutils/mdutils.py

@@ -116,29 +116,34 @@ def new_header(self, level, title, style='atx', add_table_of_contents='y'):
     def __add_new_item_table_of_content(self, level, item):
         """Automatically add new atx headers to the table of contents.
 
-        :param level: add til 2 levels. Only can take 1 or 2.
+        :param level: add levels up to 6.
         :type level: int
         :param item: items to add.
         :type item: list or str
 
         """
-        if level == 1:
-            self._table_titles.append(item)
-            self._table_titles.append([])
-        elif level == 2:
-            self._table_titles[-1].append(item)
-            self._table_titles[-1].append([])
+
+        curr = self._table_titles
+
+        for i in range(level-1):
+            curr = curr[-1]
+
+        curr.append(item)
+
+        if level < 6:
+            curr.append([])
+
 
     def new_table_of_contents(self, table_title="Table of contents", depth=1, marker=''):
         """Table of contents can be created if Headers of 'atx' style have been defined.
 
         This method allows to create a table of contents and define a title for it. Moreover, `depth` allows user to
-        define if headers of level 1 and 2 or only level 1 have to be placed on the table of contents.
+        define how many levels of headers will be placed in the table of contents.
         If no marker is defined, the table of contents will be placed automatically after the file's title.
 
         :param table_title: The table content's title, by default "Table of contents"
         :type table_title: str
-        :param depth: allows to include Headers 1 and 2 or only Headers of level 1. Possible values 1 or 2.
+        :param depth: allows to include atx headers 1 through 6. Possible values: 1, 2, 3, 4, 5, or 6.
         :type depth: int
         :param marker: allows to place the table of contents using a marker.
         :type marker: str

mdutils/tools/TableOfContents.py

@@ -16,7 +16,7 @@ def _loop_through(self, elements, tab='', depth=1):
 
         :param elements: contain all the headers defined on the main class.
         :param tab: Inserts tabulations.
-        :param depth: allows to include Headers 1 and 2 or only Headers of level 1. Possibly values 1 or 2.
+        :param depth: allows to include atx headers 1 through 6. Possible values: 1, 2, 3, 4, 5, or 6.
         :type elements: list
         :type tab: str
         :type depth: int
@@ -25,35 +25,30 @@ def _loop_through(self, elements, tab='', depth=1):
         elements_to_string = ""
         for item in elements:
             if isinstance(item, list):
-                if item and depth == 2:
-                    if tab == '\t':
-                        elements_to_string += self._loop_through(item, tab='\t\t')
-                    else:
-                        elements_to_string += self._loop_through(item, tab='\t')
+                if item and depth > 1:
+                    elements_to_string += self._loop_through(item, tab=tab+'\t', depth=depth-1)
             else:
                 elements_to_string += '\n' + tab + '* [' + item + '](#' \
                                       + re.sub('[^a-z0-9_\-]', '', item.lower().replace(' ', '-')) + ')'
 
         return elements_to_string
 
     def create_table_of_contents(self, array_of_title_contents, depth=1):
-        """This method can create a table of contents using an array of the different titles. The deep can be changed.
-        Only accepts 1 or 2.
-
+        """This method can create a table of contents using an array of the different titles. The depth can be changed.
         :param array_of_title_contents: a string list with the different headers.
         :type array_of_title_contents: list
-        :param depth: allows to include Headers 1 and 2 or only Headers of level 1. Possibly values 1 or 2.
+        :param depth: allows to include atx headers 1 through 6. Possible values: 1, 2, 3, 4, 5, or 6.
         :type depth: int
         :return: return a string ready to be written to a Markdown file.
         :rtype: str
         """
-        if depth in (1, 2):
+        if depth in range(1,7):
             table_of_contents = ""
             table_of_contents += self._loop_through(array_of_title_contents, depth=depth)
             table_of_contents += '\n'
             return table_of_contents
         else:
-            raise ValueError("depth's expected value: 1 or 2, but depth = " + str(depth))
+            raise ValueError("depth's expected value: between 1 and 6, but depth = " + str(depth))
 
 
 if __name__ == "__main__":