Markdown TOC

md-toc.py

#!/usr/bin/env python

from __future__ import unicode_literals, print_function, absolute_import, division

import os
import re
import sys
from io import StringIO, open


# ----------------------------------------------------------------------------------------------------------------------

_DEFAULT_MAX_HEAD_LEVEL = 2
# Cached re patterns as constants
_HEAD_LEVEL_PATTERN = re.compile(r"^\s*(?P<head_level>#+)\s*(?P<head_title>.*)\s*$", flags=re.MULTILINE)
_LINK_ESCAPE_PATTERN = re.compile(r"[\s']+")


def head_iter(markdown_content, max_level=_DEFAULT_MAX_HEAD_LEVEL):
    """
    Parse the content of specified markdown file to find headers

    :param str markdown_content: Content of a Markdown file
    :param int max_level: Max level of headers to yield
    :rtype: collections.Iterable[(int, str, str)]
    """
    for match in _HEAD_LEVEL_PATTERN.finditer(markdown_content):
        match_dict = match.groupdict()
        head_level = len(match_dict["head_level"])
        head_title = match_dict["head_title"]
        head_link = "#" + _LINK_ESCAPE_PATTERN.sub("-", head_title).lower()

        if head_level <= max_level:
            yield head_level, head_title, head_link


def generate_toc(markdown_content, max_level=_DEFAULT_MAX_HEAD_LEVEL):
    """
    Generate TOC for a Markdown document
    :param str markdown_content: Content of a Markdown file
    :param int max_level: Max level of headers to yield
    :rtype: str
    """
    output = StringIO()
    for head_level, head_title, head_link in head_iter(markdown_content, max_level=max_level):
        output.write("{spaces}* [{title}]({link})\n".format(
            spaces="  " * (head_level - 1),
            title=head_title,
            link=head_link
        ))
    result = output.getvalue().strip()
    output.close()
    return result


# ----------------------------------------------------------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: {0} <filename>".format(sys.argv[0]))
        exit(1)

    """Create a TOC from the specified markdown document"""
    file_path = os.path.abspath(sys.argv[1])
    if not os.path.exists(file_path):
        print("Cannot find file: {0}".format(file_path))
        exit(1)

    print("")
    with open(file_path, mode="r", encoding="utf-8") as f:
        print(generate_toc(f.read()))
    print("")