Add tools for generating contents and navigation

Author: jakevdp

.gitignore

@@ -87,3 +87,7 @@ ENV/
 
 # Rope project settings
 .ropeproject
+
+
+# Emacs
+*~

tools/add_navigation.py

@@ -0,0 +1,54 @@
+import os
+import itertools
+
+from ipykernel import kernelspec as ks
+import nbformat
+from nbformat.v4.nbbase import new_markdown_cell
+
+from generate_contents import NOTEBOOK_DIR, REG, iter_notebooks, get_notebook_title
+
+
+def prev_this_next(it):
+    a, b, c = itertools.tee(it,3)
+    next(c)
+    return zip(itertools.chain([None], a), b, itertools.chain(c, [None]))
+
+
+PREV_TEMPLATE = "< [{title}]({url}) "
+CONTENTS = "| [Contents](Index.ipynb) |"
+NEXT_TEMPLATE = " [{title}]({url}) >"
+NAV_COMMENT = "<!--NAVIGATION-->\n"
+
+
+def iter_navbars():
+    for prev_nb, nb, next_nb in prev_this_next(iter_notebooks()):
+        navbar = NAV_COMMENT
+        if prev_nb:
+            navbar += PREV_TEMPLATE.format(title=get_notebook_title(prev_nb),
+                                           url=prev_nb)
+        navbar += CONTENTS
+        if next_nb:
+            navbar += NEXT_TEMPLATE.format(title=get_notebook_title(next_nb),
+                                           url=next_nb)
+        yield os.path.join(NOTEBOOK_DIR, nb), navbar
+
+
+def write_navbars():
+    for nb_name, navbar in iter_navbars():
+        nb = nbformat.read(nb_name, as_version=4)
+        is_comment = lambda cell: cell.source.startswith(NAV_COMMENT)
+
+        if is_comment(nb.cells[1]):
+            nb.cells[1].source = navbar
+        else:
+            nb.cells.insert(1, new_markdown_cell(source=navbar))
+
+        if is_comment(nb.cells[-1]):
+            nb.cells[-1].source = navbar
+        else:
+            nb.cells.append(new_markdown_cell(source=navbar))
+        nbformat.write(nb, nb_name)
+
+
+if __name__ == '__main__':
+    write_navbars()

tools/generate_contents.py

@@ -1,10 +1,7 @@
 import os
 import re
 import itertools
-
-PREV_TEMPLATE = " <[{title}]({url}) "
-CONTENTS = "| [Contents](Index.ipynb)| "
-NEXT_TEMPLATE = " [{title}](url) >"
+import nbformat
 
 NOTEBOOK_DIR = os.path.join(os.path.dirname(__file__), '..', 'notebooks')
 
@@ -17,40 +14,37 @@
 
 REG = re.compile(r'(\d\d)\.(\d\d)-(.*)\.ipynb')
 
-notebooks = sorted(nb for nb in os.listdir(NOTEBOOK_DIR) if REG.match(nb))
-
-def prev_this_next(it):
-    a, b, c = itertools.tee(it,3)
-    next(c)
-    return zip(itertools.chain([None], a), b, itertools.chain(c, [None]))
-
 
+def iter_notebooks():
+    return sorted(nb for nb in os.listdir(NOTEBOOK_DIR) if REG.match(nb))
 
 
-def iter_navbars(notebooks):
-    for prev_nb, nb, next_nb in prev_this_next(notebooks):
-        navbar = ""
-        if prev_nb:
-            navbar += PREV_TEMPLATE.format(title=REG.match(prev_nb).groups()[2],
-                                           url=prev_nb)
-        navbar += CONTENTS
-        if next_nb:
-            navbar += NEXT_TEMPLATE.format(title=REG.match(next_nb).groups()[2],
-                                           url=next_nb)
-        yield navbar
+def get_notebook_title(nb_file):
+    nb = nbformat.read(os.path.join(NOTEBOOK_DIR, nb_file), as_version=4)
+    for cell in nb.cells:
+        if cell.source.startswith('#'):
+            return cell.source[1:].splitlines()[0].strip()
 
 
-def gen_contents(notebooks):
-    def get_chapter(nb):
-        return REG.match(nb).groups()[0]
-    
-    for nb in notebooks:
+def gen_contents(directory=None):
+    for nb in iter_notebooks():
+        if directory:
+            nb_url = os.path.join(directory, nb)
+        else:
+            nb_url = nb
         chapter, section, title = REG.match(nb).groups()
-        title = title.replace('-', ' ')
+        title = get_notebook_title(nb)
         if section == '00':
-            yield '\n### [{0}]({1})'.format(title, nb)
+            yield '\n### [{0}]({1})'.format(title, nb_url)
         else:
-            yield "- [{0}]({1})".format(title, nb)
+            yield "- [{0}]({1})".format(title, nb_url)
+
+
+def print_contents(directory=None):
+    print('\n'.join(gen_contents(directory)))
 
 
-print('\n'.join(gen_contents(notebooks)))
+if __name__ == '__main__':
+    print_contents()
+    print(70 * '#')
+    print_contents('notebooks')