Create a Table of Contents for an iPython Notebook.

ipynb_toc.py

import json
import re
import string
from IPython.display import display_markdown

class TOC:
    def __init__(self, notebook_path):
        """
        Class that builds a table of content for a given notebook. For now, it
        only accepts markdown headers using the '#' format (e.g.: ## This is a
        header of level 2). You can use it directly in your notebook as follows:
        '''
        %run ipynb_toc.py
        TOC('mynotebook.ipynb').display_markdown(min_level=2)
        '''
        """
        self.notebook_path = notebook_path
        self.json_notebook = None
        self.toc = []
        self.title_regex = re.compile('\#+')

    def get_toc_content(self, cell):
        """
        Append the potential headers contained in this cell to the toc list.
        """
        for line in cell['source']:
            m = self.title_regex.match(line)
            if m:
                level = len(m.group()) # count number of # character
                title = line.strip('#').strip()
                href = title.replace(' ', '-')
                self.toc.append((level, title, href))

    def create_toc(self):
        """
        Iterate through all the markdown cells in the notebook and fill a list
        containing the headers.
        """
        with open(self.notebook_path) as nb_fd:
            self.json_notebook = json.load(nb_fd)
        self.markdown_cells = [cell for cell in self.json_notebook['cells'] if cell['cell_type'] == 'markdown']
        for mcell in self.markdown_cells:
            self.get_toc_content(mcell)

    def display_markdown(self, min_level=1):
        """
        Create the table of contents and display it in markdown.
        """
        md_content = self.to_markdown(min_level)
        display_markdown(md_content, raw=True)

    def to_markdown(self, min_level=1):
        """
        Create the table and return it in markdown.
        """
        self.create_toc()
        header_list = []
        href_list = []
        level_count = 1
        previous_level = 0
        for idx, (level, title, href) in enumerate(self.toc):
            if level >= min_level:
                if previous_level == level:
                    level_count += 1
                else:
                    level_count = 1
                header = '%s%d. [%s][%d]' % ('  '*(level-min_level),
                                             level_count, title, idx)
                header_list.append(header)
                href_list.append('[%d]: #%s' % (idx, href))
            previous_level = level
        md_content = "\n".join(header_list + href_list)
        return md_content