/
postproc.py
70 lines (58 loc) · 2.35 KB
/
postproc.py
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
#! /usr/bin/env python3
from pathlib import Path
import nbformat
MAGIC_PREFIX = "##"
def load_notebook(nbfile):
with open(nbfile, "r") as f:
return nbformat.read(f, nbformat.NO_CONVERT)
def save_notebook(nb, nbfile):
with open(nbfile, "w") as f:
nbformat.write(nb, f)
def process_notebook(nb):
nb_nosol = nbformat.v4.new_notebook()
nb_sol = nbformat.v4.new_notebook()
txt_sol = []
for nnb in (nb_nosol, nb_sol):
nnb.metadata.language_info = nb.metadata.language_info
for cell in nb.cells:
write_to = [nb_nosol, nb_sol]
if cell.cell_type == "code" and cell.source.startswith(MAGIC_PREFIX):
magic, _, cell.source = cell.source.partition("\n")
magic = magic[len(MAGIC_PREFIX):].split(":")
if magic[0] == "sol":
write_to = [nb_sol]
if len(magic) > 1:
txt_sol.append("# {}\n{}".format(magic[1], cell.source))
else:
txt_sol.append(cell.source)
elif magic[0] == "nosol":
write_to = [nb_nosol]
elif magic[0] == "solhead":
write_to = []
txt_sol.append("{}# {}".format("\n" if txt_sol else "", magic[1]))
elif magic[0] == "outsnip":
write_to = [nb_nosol, nb_sol]
for cout in cell.outputs:
if cout.name == "stdout" and cout.output_type == "stream":
lines = cout.text.split("\n")
new_lines = lines[:int(magic[1])] + ["[...]"] + lines[-int(magic[2]):]
cout.text = "\n".join(new_lines)
for dnb in write_to:
dnb.cells.append(cell)
return nb_nosol, nb_sol, "\n\n".join(txt_sol)
def main():
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument("notebook_file")
args = parser.parse_args()
input_path = Path(args.notebook_file)
base_file_name = input_path.stem
nosol_path = input_path.parent / (base_file_name + ".nosol.ipynb")
sol_path = input_path.parent / (base_file_name + ".sol.ipynb")
input_nb = load_notebook(input_path)
nosol_nb, sol_nb, sol_txt = process_notebook(input_nb)
save_notebook(nosol_nb, nosol_path)
save_notebook(sol_nb, sol_path)
print(sol_txt)
if __name__ == "__main__":
main()