Coverage for muutils\nbutils\run_notebook_tests.py: 65%
74 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-10-09 01:48 -0600
« prev ^ index » next coverage.py v7.6.1, created at 2024-10-09 01:48 -0600
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
15from muutils.spinner import SpinnerContext
18class NotebookTestError(Exception):
19 pass
22def run_notebook_tests(
23 notebooks_dir: Path,
24 converted_notebooks_temp_dir: Path,
25 CI_output_suffix: str = ".CI-output.txt",
26 run_python_cmd: str = "poetry run python",
27 exit_on_first_fail: bool = False,
28):
29 original_cwd: Path = Path.cwd()
30 # get paths
31 notebooks_dir = Path(notebooks_dir)
32 converted_notebooks_temp_dir = Path(converted_notebooks_temp_dir)
33 root_relative_to_notebooks: Path = Path(os.path.relpath(".", notebooks_dir))
35 term_width: int
36 try:
37 term_width = os.get_terminal_size().columns
38 except OSError:
39 term_width = 80
41 exceptions: dict[str, str] = dict()
43 print(f"# testing notebooks in '{notebooks_dir}'")
44 print(
45 f"# reading converted notebooks from '{converted_notebooks_temp_dir.as_posix()}'"
46 )
48 try:
49 # check things exist
50 if not notebooks_dir.exists():
51 raise NotebookTestError(f"Notebooks dir '{notebooks_dir}' does not exist")
52 if not notebooks_dir.is_dir():
53 raise NotebookTestError(
54 f"Notebooks dir '{notebooks_dir}' is not a directory"
55 )
56 if not converted_notebooks_temp_dir.exists():
57 raise NotebookTestError(
58 f"Converted notebooks dir '{converted_notebooks_temp_dir}' does not exist"
59 )
60 if not converted_notebooks_temp_dir.is_dir():
61 raise NotebookTestError(
62 f"Converted notebooks dir '{converted_notebooks_temp_dir}' is not a directory"
63 )
65 notebooks: list[Path] = list(notebooks_dir.glob("*.ipynb"))
66 if not notebooks:
67 raise NotebookTestError(f"No notebooks found in '{notebooks_dir}'")
69 converted_notebooks: list[Path] = list()
70 for nb in notebooks:
71 converted_file: Path = (
72 converted_notebooks_temp_dir / nb.with_suffix(".py").name
73 )
74 if not converted_file.exists():
75 raise NotebookTestError(
76 f"Did not find converted notebook '{converted_file}' for '{nb}'"
77 )
78 converted_notebooks.append(converted_file)
80 del converted_file
82 # the location of this line is important
83 os.chdir(notebooks_dir)
85 n_notebooks: int = len(converted_notebooks)
86 for idx, file in enumerate(converted_notebooks):
87 # run the file
88 print(f"Running {idx+1}/{n_notebooks}: {file.as_posix()}")
89 output_file: Path = file.with_suffix(CI_output_suffix)
90 print(f" Output in {output_file.as_posix()}")
91 with SpinnerContext(
92 spinner_chars="braille",
93 update_interval=0.5,
94 format_string="\r {spinner} ({elapsed_time:.2f}s) {message}{value}",
95 ):
96 command: str = f"{run_python_cmd} {root_relative_to_notebooks / file} > {root_relative_to_notebooks / output_file} 2>&1"
97 process: subprocess.CompletedProcess = subprocess.run(
98 command, shell=True, text=True
99 )
101 if process.returncode == 0:
102 print(f" ✅ Run completed with return code {process.returncode}")
103 else:
104 print(
105 f" ❌ Run failed with return code {process.returncode}!!! Check {output_file.as_posix()}"
106 )
108 # print the output of the file to the console if it failed
109 if process.returncode != 0:
110 with open(root_relative_to_notebooks / output_file, "r") as f:
111 file_output: str = f.read()
112 err: str = f"Error in {file}:\n{'-'*term_width}\n{file_output}"
113 exceptions[file.as_posix()] = err
114 if exit_on_first_fail:
115 raise NotebookTestError(err)
117 del process
119 if len(exceptions) > 0:
120 exceptions_str: str = ("\n" + "=" * term_width + "\n").join(
121 list(exceptions.values())
122 )
123 raise NotebookTestError(
124 exceptions_str
125 + "=" * term_width
126 + f"\n❌ {len(exceptions)}/{n_notebooks} notebooks failed:\n{list(exceptions.keys())}"
127 )
129 except NotebookTestError as e:
130 print("!" * term_width, file=sys.stderr)
131 print(e, file=sys.stderr)
132 print("!" * term_width, file=sys.stderr)
133 raise e
134 finally:
135 # return to original cwd
136 os.chdir(original_cwd)
139if __name__ == "__main__":
140 import argparse
142 parser: argparse.ArgumentParser = argparse.ArgumentParser()
144 parser.add_argument(
145 "--notebooks-dir",
146 type=str,
147 help="The directory from which to run the notebooks",
148 )
149 parser.add_argument(
150 "--converted-notebooks-temp-dir",
151 type=str,
152 help="The directory containing the converted notebooks to test",
153 )
155 args: argparse.Namespace = parser.parse_args()
157 run_notebook_tests(
158 Path(args.notebooks_dir),
159 Path(args.converted_notebooks_temp_dir),
160 )