forked from empa-scientific-it/python-tutorial
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtoc.py
More file actions
executable file
·102 lines (78 loc) · 2.96 KB
/
Copy pathtoc.py
File metadata and controls
executable file
·102 lines (78 loc) · 2.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
#!/usr/bin/env python
"""CLI script to build a table of contents for an IPython notebook"""
import argparse as ap
import pathlib
import re
from collections import namedtuple
import nbformat
from nbformat import NotebookNode
TocEntry = namedtuple("TocEntry", ["level", "text", "anchor"])
def extract_markdown_cells(notebook: NotebookNode) -> str:
"""Extract the markdown cells from a notebook"""
return "\n".join(
[cell.source for cell in notebook.cells if cell.cell_type == "markdown"]
)
def extract_toc(notebook: str) -> list[TocEntry]:
"""Extract the table of contents from a markdown string"""
toc = []
line_re = re.compile(r"(#+)\s+(.+)")
for line in notebook.splitlines():
if groups := re.match(line_re, line):
heading, text, *_ = groups.groups()
level = len(heading)
anchor = "-".join(text.replace("`", "").split())
toc.append(TocEntry(level, text, anchor))
return toc
def markdown_toc(toc: list[TocEntry]) -> str:
"""Build a string representation of the toc as a nested markdown list"""
lines = []
for entry in toc:
line = f"{' ' * entry.level}- [{entry.text}](#{entry.anchor})"
lines.append(line)
return "\n".join(lines)
def build_toc(nb_path: pathlib.Path, placeholder: str = "[TOC]") -> NotebookNode:
"""Build a table of contents for a notebook and insert it at the location of a placeholder"""
# Read the notebook
nb_obj: NotebookNode = nbformat.read(nb_path, nbformat.NO_CONVERT)
md_cells = extract_markdown_cells(nb_obj)
# Build tree
toc_tree = extract_toc(md_cells)
# Build toc representation
toc_repr = markdown_toc(toc_tree)
# Insert it a the location of a placeholder
toc_header = "# Table of Contents"
for cell in nb_obj.cells:
if cell.source.startswith((placeholder, toc_header)):
cell.source = f"{toc_header}\n{toc_repr}"
cell.cell_type = "markdown"
return nb_obj
def main():
"""CLI entry point"""
parser = ap.ArgumentParser(
description="Build a table of contents for an IPython notebook"
)
parser.add_argument("notebook", type=str, help="Path to the notebook to process")
parser.add_argument(
"--output", "-o", type=str, default=None, help="Path to the output notebook"
)
parser.add_argument(
"--force",
"-f",
action="store_true",
default=False,
help="Force overwrite of original notebook",
)
args = parser.parse_args()
if not (input_nb := pathlib.Path(args.notebook)).exists():
raise FileNotFoundError(input_nb)
if args.output is None:
output_nb = input_nb.with_suffix(".toc.ipynb")
else:
output_nb = pathlib.Path(args.output)
with output_nb.open("w", encoding="utf-8") as file:
nbformat.write(build_toc(input_nb), file)
if args.force:
input_nb.unlink()
output_nb.rename(input_nb)
if __name__ == "__main__":
main()