Coverage for muutils\nbutils\run_notebook_tests.py: 61%
85 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-01-17 00:06 -0700
« prev ^ index » next coverage.py v7.6.1, created at 2025-01-17 00:06 -0700
1"""turn a folder of notebooks into scripts, run them, and make sure they work.
3made to be called as
5```bash
6python -m muutils.nbutils.run_notebook_tests --notebooks-dir <notebooks_dir> --converted-notebooks-temp-dir <converted_notebooks_temp_dir>
7```
8"""
10import os
11import subprocess
12import sys
13from pathlib import Path
14from typing import Optional
15import warnings
17from muutils.spinner import SpinnerContext
20class NotebookTestError(Exception):
21 pass
24def run_notebook_tests(
25 notebooks_dir: Path,
26 converted_notebooks_temp_dir: Path,
27 CI_output_suffix: str = ".CI-output.txt",
28 run_python_cmd: Optional[str] = None,
29 run_python_cmd_fmt: str = "{python_tool} run python",
30 python_tool: str = "poetry",
31 exit_on_first_fail: bool = False,
32):
33 """Run converted Jupyter notebooks as Python scripts and verify they execute successfully.
35 Takes a directory of notebooks and their corresponding converted Python scripts,
36 executes each script, and captures the output. Failures are collected and reported,
37 with optional early exit on first failure.
39 # Parameters:
40 - `notebooks_dir : Path`
41 Directory containing the original .ipynb notebook files
42 - `converted_notebooks_temp_dir : Path`
43 Directory containing the corresponding converted .py files
44 - `CI_output_suffix : str`
45 Suffix to append to output files capturing execution results
46 (defaults to `".CI-output.txt"`)
47 - `run_python_cmd : str | None`
48 Custom command to run Python scripts. Overrides python_tool and run_python_cmd_fmt if provided
49 (defaults to `None`)
50 - `run_python_cmd_fmt : str`
51 Format string for constructing the Python run command
52 (defaults to `"{python_tool} run python"`)
53 - `python_tool : str`
54 Tool used to run Python (e.g. poetry, uv)
55 (defaults to `"poetry"`)
56 - `exit_on_first_fail : bool`
57 Whether to raise exception immediately on first notebook failure
58 (defaults to `False`)
60 # Returns:
61 - `None`
63 # Modifies:
64 - Working directory: Temporarily changes to notebooks_dir during execution
65 - Filesystem: Creates output files with CI_output_suffix for each notebook
67 # Raises:
68 - `NotebookTestError`: If any notebooks fail to execute, or if input directories are invalid
69 - `TypeError`: If run_python_cmd is provided but not a string
71 # Usage:
72 ```python
73 >>> run_notebook_tests(
74 ... notebooks_dir=Path("notebooks"),
75 ... converted_notebooks_temp_dir=Path("temp/converted"),
76 ... python_tool="poetry"
77 ... )
78 # testing notebooks in 'notebooks'
79 # reading converted notebooks from 'temp/converted'
80 Running 1/2: temp/converted/notebook1.py
81 Output in temp/converted/notebook1.CI-output.txt
82 ✅ Run completed with return code 0
83 ```
84 """
86 run_python_cmd_: str
87 if run_python_cmd is None:
88 run_python_cmd_ = run_python_cmd_fmt.format(python_tool=python_tool)
89 elif isinstance(run_python_cmd, str):
90 run_python_cmd_ = run_python_cmd
91 warnings.warn(
92 "You have specified a custom run_python_cmd, this will override the `python_tool` parameter and `run_python_cmd_fmt` parameter. This will be removed in a future version.",
93 DeprecationWarning,
94 )
95 else:
96 raise TypeError(
97 f"run_python_cmd must be a string or None, got {run_python_cmd =}, {type(run_python_cmd) =}"
98 )
100 original_cwd: Path = Path.cwd()
101 # get paths
102 notebooks_dir = Path(notebooks_dir)
103 converted_notebooks_temp_dir = Path(converted_notebooks_temp_dir)
104 root_relative_to_notebooks: Path = Path(os.path.relpath(".", notebooks_dir))
106 term_width: int
107 try:
108 term_width = os.get_terminal_size().columns
109 except OSError:
110 term_width = 80
112 exceptions: dict[str, str] = dict()
114 print(f"# testing notebooks in '{notebooks_dir}'")
115 print(
116 f"# reading converted notebooks from '{converted_notebooks_temp_dir.as_posix()}'"
117 )
119 try:
120 # check things exist
121 if not notebooks_dir.exists():
122 raise NotebookTestError(f"Notebooks dir '{notebooks_dir}' does not exist")
123 if not notebooks_dir.is_dir():
124 raise NotebookTestError(
125 f"Notebooks dir '{notebooks_dir}' is not a directory"
126 )
127 if not converted_notebooks_temp_dir.exists():
128 raise NotebookTestError(
129 f"Converted notebooks dir '{converted_notebooks_temp_dir}' does not exist"
130 )
131 if not converted_notebooks_temp_dir.is_dir():
132 raise NotebookTestError(
133 f"Converted notebooks dir '{converted_notebooks_temp_dir}' is not a directory"
134 )
136 notebooks: list[Path] = list(notebooks_dir.glob("*.ipynb"))
137 if not notebooks:
138 raise NotebookTestError(f"No notebooks found in '{notebooks_dir}'")
140 converted_notebooks: list[Path] = list()
141 for nb in notebooks:
142 converted_file: Path = (
143 converted_notebooks_temp_dir / nb.with_suffix(".py").name
144 )
145 if not converted_file.exists():
146 raise NotebookTestError(
147 f"Did not find converted notebook '{converted_file}' for '{nb}'"
148 )
149 converted_notebooks.append(converted_file)
151 del converted_file
153 # the location of this line is important
154 os.chdir(notebooks_dir)
156 n_notebooks: int = len(converted_notebooks)
157 for idx, file in enumerate(converted_notebooks):
158 # run the file
159 print(f"Running {idx+1}/{n_notebooks}: {file.as_posix()}")
160 output_file: Path = file.with_suffix(CI_output_suffix)
161 print(f" Output in {output_file.as_posix()}")
162 with SpinnerContext(
163 spinner_chars="braille",
164 update_interval=0.5,
165 format_string="\r {spinner} ({elapsed_time:.2f}s) {message}{value}",
166 ):
167 command: str = f"{run_python_cmd_} {root_relative_to_notebooks / file} > {root_relative_to_notebooks / output_file} 2>&1"
168 process: subprocess.CompletedProcess = subprocess.run(
169 command, shell=True, text=True
170 )
172 if process.returncode == 0:
173 print(f" ✅ Run completed with return code {process.returncode}")
174 else:
175 print(
176 f" ❌ Run failed with return code {process.returncode}!!! Check {output_file.as_posix()}"
177 )
179 # print the output of the file to the console if it failed
180 if process.returncode != 0:
181 with open(root_relative_to_notebooks / output_file, "r") as f:
182 file_output: str = f.read()
183 err: str = f"Error in {file}:\n{'-'*term_width}\n{file_output}"
184 exceptions[file.as_posix()] = err
185 if exit_on_first_fail:
186 raise NotebookTestError(err)
188 del process
190 if len(exceptions) > 0:
191 exceptions_str: str = ("\n" + "=" * term_width + "\n").join(
192 list(exceptions.values())
193 )
194 raise NotebookTestError(
195 exceptions_str
196 + "=" * term_width
197 + f"\n❌ {len(exceptions)}/{n_notebooks} notebooks failed:\n{list(exceptions.keys())}"
198 )
200 except NotebookTestError as e:
201 print("!" * term_width, file=sys.stderr)
202 print(e, file=sys.stderr)
203 print("!" * term_width, file=sys.stderr)
204 raise e
205 finally:
206 # return to original cwd
207 os.chdir(original_cwd)
210if __name__ == "__main__":
211 import argparse
213 parser: argparse.ArgumentParser = argparse.ArgumentParser()
215 parser.add_argument(
216 "--notebooks-dir",
217 type=str,
218 help="The directory from which to run the notebooks",
219 )
220 parser.add_argument(
221 "--converted-notebooks-temp-dir",
222 type=str,
223 help="The directory containing the converted notebooks to test",
224 )
225 parser.add_argument(
226 "--python-tool",
227 type=str,
228 default="poetry",
229 help="The python tool to use to run the notebooks (usually uv or poetry)",
230 )
231 parser.add_argument(
232 "--run-python-cmd-fmt",
233 type=str,
234 default="{python_tool} run python",
235 help="The command to run python with the python tool. if you don't want to use poetry or uv, you can just set this to 'python'",
236 )
237 parser.add_argument(
238 "--help",
239 "-h",
240 action="help",
241 help="Show this help message and exit",
242 )
244 args: argparse.Namespace = parser.parse_args()
246 run_notebook_tests(
247 notebooks_dir=Path(args.notebooks_dir),
248 converted_notebooks_temp_dir=Path(args.converted_notebooks_temp_dir),
249 python_tool=args.python_tool,
250 run_python_cmd_fmt=args.run_python_cmd_fmt,
251 )