1+ import click
2+ import json
3+ import os
4+
5+ DOCS_PATH = os .path .join (os .path .dirname (os .path .abspath (__file__ )), "source" )
6+
7+
8+ def _get_ipython_notebooks (docs_source ):
9+ directories_to_skip = ["_templates" , "generated" , ".ipynb_checkpoints" ]
10+ notebooks = []
11+ for root , _ , filenames in os .walk (docs_source ):
12+ if any (dir_ in root for dir_ in directories_to_skip ):
13+ continue
14+ for filename in filenames :
15+ if filename .endswith (".ipynb" ):
16+ notebooks .append (os .path .join (root , filename ))
17+ return notebooks
18+
19+
20+ def _check_delete_empty_cell (notebook , delete = True ):
21+ with open (notebook , "r" ) as f :
22+ source = json .load (f )
23+ cell = source ["cells" ][- 1 ]
24+ if cell ["cell_type" ] == "code" and cell ["source" ] == []:
25+ # this is an empty cell, which we should delete
26+ if delete :
27+ source ["cells" ] = source ["cells" ][:- 1 ]
28+ else :
29+ return False
30+ if delete :
31+ with open (notebook , "w" ) as f :
32+ json .dump (source , f , ensure_ascii = False , indent = 1 )
33+ else :
34+ return True
35+
36+
37+ def _check_execution_and_output (notebook ):
38+ with open (notebook , "r" ) as f :
39+ source = json .load (f )
40+ for cells in source ["cells" ]:
41+ if cells ["cell_type" ] == "code" and (cells ["execution_count" ] is not None or cells ['outputs' ]!= []):
42+ return False
43+ return True
44+
45+
46+ def _check_python_version (notebook , default_version ):
47+ with open (notebook , "r" ) as f :
48+ source = json .load (f )
49+ if source ["metadata" ]["language_info" ]["version" ] != default_version :
50+ return False
51+ return True
52+
53+
54+ def _fix_python_version (notebook , default_version ):
55+ with open (notebook , "r" ) as f :
56+ source = json .load (f )
57+ source ["metadata" ]["language_info" ]["version" ] = default_version
58+ with open (notebook , "w" ) as f :
59+ json .dump (source , f , ensure_ascii = False , indent = 1 )
60+
61+
62+ def _fix_execution_and_output (notebook ):
63+ with open (notebook , "r" ) as f :
64+ source = json .load (f )
65+ for cells in source ["cells" ]:
66+ if cells ["cell_type" ] == "code" and cells ["execution_count" ] is not None :
67+ cells ["execution_count" ] = None
68+ cells ["outputs" ] = []
69+ source ["metadata" ]["kernelspec" ]["display_name" ] = "Python 3"
70+ source ["metadata" ]["kernelspec" ]["name" ] = "python3"
71+ with open (notebook , "w" ) as f :
72+ json .dump (source , f , ensure_ascii = False , indent = 1 )
73+
74+
75+ def _get_notebooks_with_executions_and_empty (notebooks , default_version = "3.8.2" ):
76+ executed = []
77+ empty_last_cell = []
78+ versions = []
79+ for notebook in notebooks :
80+ if not _check_execution_and_output (notebook ):
81+ executed .append (notebook )
82+ if not _check_delete_empty_cell (notebook , delete = False ):
83+ empty_last_cell .append (notebook )
84+ if not _check_python_version (notebook , default_version ):
85+ versions .append (notebook )
86+ return (executed , empty_last_cell , versions )
87+
88+
89+ def _fix_versions (notebooks , default_version = "3.8.2" ):
90+ for notebook in notebooks :
91+ _fix_python_version (notebook , default_version )
92+
93+
94+ def _remove_notebook_empty_last_cell (notebooks ):
95+ for notebook in notebooks :
96+ _check_delete_empty_cell (notebook , delete = True )
97+
98+
99+ def _standardize_outputs (notebooks ):
100+ for notebook in notebooks :
101+ _fix_execution_and_output (notebook )
102+
103+
104+ @click .group ()
105+ def cli ():
106+ """no-op"""
107+
108+
109+ @cli .command ()
110+ def standardize ():
111+ notebooks = _get_ipython_notebooks (DOCS_PATH )
112+ executed_notebooks , empty_cells , versions = _get_notebooks_with_executions_and_empty (notebooks )
113+ if executed_notebooks :
114+ _standardize_outputs (executed_notebooks )
115+ executed_notebooks = ["\t " + notebook for notebook in executed_notebooks ]
116+ executed_notebooks = "\n " .join (executed_notebooks )
117+ click .echo (
118+ f"Removed the outputs for:\n { executed_notebooks } "
119+ )
120+ if empty_cells :
121+ _remove_notebook_empty_last_cell (empty_cells )
122+ empty_cells = ["\t " + notebook for notebook in empty_cells ]
123+ empty_cells = "\n " .join (empty_cells )
124+ click .echo (
125+ f"Removed the empty cells for:\n { empty_cells } "
126+ )
127+ if versions :
128+ _fix_versions (versions )
129+ versions = ["\t " + notebook for notebook in versions ]
130+ versions = "\n " .join (versions )
131+ click .echo (
132+ f"Fixed python versions for:\n { versions } "
133+ )
134+
135+
136+ @cli .command ()
137+ def check_execution ():
138+ notebooks = _get_ipython_notebooks (DOCS_PATH )
139+ executed_notebooks , empty_cells , versions = _get_notebooks_with_executions_and_empty (notebooks )
140+ if executed_notebooks :
141+ executed_notebooks = ["\t " + notebook for notebook in executed_notebooks ]
142+ executed_notebooks = "\n " .join (executed_notebooks )
143+ raise SystemExit (
144+ f"The following notebooks have executed outputs:\n { executed_notebooks } \n "
145+ "Please run make lint-fix to fix this."
146+ )
147+ if empty_cells :
148+ empty_cells = ["\t " + notebook for notebook in empty_cells ]
149+ empty_cells = "\n " .join (empty_cells )
150+ raise SystemExit (
151+ f"The following notebooks have empty cells at the end:\n { empty_cells } \n "
152+ "Please run make lint-fix to fix this."
153+ )
154+ if versions :
155+ versions = ["\t " + notebook for notebook in versions ]
156+ versions = "\n " .join (versions )
157+ raise SystemExit (
158+ f"The following notebooks have the wrong Python version: \n { versions } \n "
159+ "Please run make lint-fix to fix this."
160+ )
161+
162+ if __name__ == "__main__" :
163+ cli ()
0 commit comments