muutils.nbutils.run_notebook_tests
turn a folder of notebooks into scripts, run them, and make sure they work.
made to be called as
python -m muutils.nbutils.run_notebook_tests --notebooks-dir <notebooks_dir> --converted-notebooks-temp-dir <converted_notebooks_temp_dir>
1"""turn a folder of notebooks into scripts, run them, and make sure they work. 2 3made to be called as 4 5```bash 6python -m muutils.nbutils.run_notebook_tests --notebooks-dir <notebooks_dir> --converted-notebooks-temp-dir <converted_notebooks_temp_dir> 7``` 8""" 9 10import os 11import subprocess 12import sys 13from pathlib import Path 14from typing import Optional 15import warnings 16 17from muutils.spinner import SpinnerContext 18 19 20class NotebookTestError(Exception): 21 pass 22 23 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. 34 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. 38 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`) 59 60 # Returns: 61 - `None` 62 63 # Modifies: 64 - Working directory: Temporarily changes to notebooks_dir during execution 65 - Filesystem: Creates output files with CI_output_suffix for each notebook 66 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 70 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 """ 85 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 ) 99 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)) 105 106 term_width: int 107 try: 108 term_width = os.get_terminal_size().columns 109 except OSError: 110 term_width = 80 111 112 exceptions: dict[str, str] = dict() 113 114 print(f"# testing notebooks in '{notebooks_dir}'") 115 print( 116 f"# reading converted notebooks from '{converted_notebooks_temp_dir.as_posix()}'" 117 ) 118 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 ) 135 136 notebooks: list[Path] = list(notebooks_dir.glob("*.ipynb")) 137 if not notebooks: 138 raise NotebookTestError(f"No notebooks found in '{notebooks_dir}'") 139 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) 150 151 del converted_file 152 153 # the location of this line is important 154 os.chdir(notebooks_dir) 155 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 ) 171 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 ) 178 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) 187 188 del process 189 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 ) 199 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) 208 209 210if __name__ == "__main__": 211 import argparse 212 213 parser: argparse.ArgumentParser = argparse.ArgumentParser() 214 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 ) 243 244 args: argparse.Namespace = parser.parse_args() 245 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 )
class
NotebookTestError(builtins.Exception):
Common base class for all non-exit exceptions.
Inherited Members
- builtins.Exception
- Exception
- builtins.BaseException
- with_traceback
- add_note
- args
def
run_notebook_tests( notebooks_dir: pathlib.Path, converted_notebooks_temp_dir: pathlib.Path, CI_output_suffix: str = '.CI-output.txt', run_python_cmd: Optional[str] = None, run_python_cmd_fmt: str = '{python_tool} run python', python_tool: str = 'poetry', exit_on_first_fail: bool = False):
25def run_notebook_tests( 26 notebooks_dir: Path, 27 converted_notebooks_temp_dir: Path, 28 CI_output_suffix: str = ".CI-output.txt", 29 run_python_cmd: Optional[str] = None, 30 run_python_cmd_fmt: str = "{python_tool} run python", 31 python_tool: str = "poetry", 32 exit_on_first_fail: bool = False, 33): 34 """Run converted Jupyter notebooks as Python scripts and verify they execute successfully. 35 36 Takes a directory of notebooks and their corresponding converted Python scripts, 37 executes each script, and captures the output. Failures are collected and reported, 38 with optional early exit on first failure. 39 40 # Parameters: 41 - `notebooks_dir : Path` 42 Directory containing the original .ipynb notebook files 43 - `converted_notebooks_temp_dir : Path` 44 Directory containing the corresponding converted .py files 45 - `CI_output_suffix : str` 46 Suffix to append to output files capturing execution results 47 (defaults to `".CI-output.txt"`) 48 - `run_python_cmd : str | None` 49 Custom command to run Python scripts. Overrides python_tool and run_python_cmd_fmt if provided 50 (defaults to `None`) 51 - `run_python_cmd_fmt : str` 52 Format string for constructing the Python run command 53 (defaults to `"{python_tool} run python"`) 54 - `python_tool : str` 55 Tool used to run Python (e.g. poetry, uv) 56 (defaults to `"poetry"`) 57 - `exit_on_first_fail : bool` 58 Whether to raise exception immediately on first notebook failure 59 (defaults to `False`) 60 61 # Returns: 62 - `None` 63 64 # Modifies: 65 - Working directory: Temporarily changes to notebooks_dir during execution 66 - Filesystem: Creates output files with CI_output_suffix for each notebook 67 68 # Raises: 69 - `NotebookTestError`: If any notebooks fail to execute, or if input directories are invalid 70 - `TypeError`: If run_python_cmd is provided but not a string 71 72 # Usage: 73 ```python 74 >>> run_notebook_tests( 75 ... notebooks_dir=Path("notebooks"), 76 ... converted_notebooks_temp_dir=Path("temp/converted"), 77 ... python_tool="poetry" 78 ... ) 79 # testing notebooks in 'notebooks' 80 # reading converted notebooks from 'temp/converted' 81 Running 1/2: temp/converted/notebook1.py 82 Output in temp/converted/notebook1.CI-output.txt 83 ✅ Run completed with return code 0 84 ``` 85 """ 86 87 run_python_cmd_: str 88 if run_python_cmd is None: 89 run_python_cmd_ = run_python_cmd_fmt.format(python_tool=python_tool) 90 elif isinstance(run_python_cmd, str): 91 run_python_cmd_ = run_python_cmd 92 warnings.warn( 93 "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.", 94 DeprecationWarning, 95 ) 96 else: 97 raise TypeError( 98 f"run_python_cmd must be a string or None, got {run_python_cmd =}, {type(run_python_cmd) =}" 99 ) 100 101 original_cwd: Path = Path.cwd() 102 # get paths 103 notebooks_dir = Path(notebooks_dir) 104 converted_notebooks_temp_dir = Path(converted_notebooks_temp_dir) 105 root_relative_to_notebooks: Path = Path(os.path.relpath(".", notebooks_dir)) 106 107 term_width: int 108 try: 109 term_width = os.get_terminal_size().columns 110 except OSError: 111 term_width = 80 112 113 exceptions: dict[str, str] = dict() 114 115 print(f"# testing notebooks in '{notebooks_dir}'") 116 print( 117 f"# reading converted notebooks from '{converted_notebooks_temp_dir.as_posix()}'" 118 ) 119 120 try: 121 # check things exist 122 if not notebooks_dir.exists(): 123 raise NotebookTestError(f"Notebooks dir '{notebooks_dir}' does not exist") 124 if not notebooks_dir.is_dir(): 125 raise NotebookTestError( 126 f"Notebooks dir '{notebooks_dir}' is not a directory" 127 ) 128 if not converted_notebooks_temp_dir.exists(): 129 raise NotebookTestError( 130 f"Converted notebooks dir '{converted_notebooks_temp_dir}' does not exist" 131 ) 132 if not converted_notebooks_temp_dir.is_dir(): 133 raise NotebookTestError( 134 f"Converted notebooks dir '{converted_notebooks_temp_dir}' is not a directory" 135 ) 136 137 notebooks: list[Path] = list(notebooks_dir.glob("*.ipynb")) 138 if not notebooks: 139 raise NotebookTestError(f"No notebooks found in '{notebooks_dir}'") 140 141 converted_notebooks: list[Path] = list() 142 for nb in notebooks: 143 converted_file: Path = ( 144 converted_notebooks_temp_dir / nb.with_suffix(".py").name 145 ) 146 if not converted_file.exists(): 147 raise NotebookTestError( 148 f"Did not find converted notebook '{converted_file}' for '{nb}'" 149 ) 150 converted_notebooks.append(converted_file) 151 152 del converted_file 153 154 # the location of this line is important 155 os.chdir(notebooks_dir) 156 157 n_notebooks: int = len(converted_notebooks) 158 for idx, file in enumerate(converted_notebooks): 159 # run the file 160 print(f"Running {idx+1}/{n_notebooks}: {file.as_posix()}") 161 output_file: Path = file.with_suffix(CI_output_suffix) 162 print(f" Output in {output_file.as_posix()}") 163 with SpinnerContext( 164 spinner_chars="braille", 165 update_interval=0.5, 166 format_string="\r {spinner} ({elapsed_time:.2f}s) {message}{value}", 167 ): 168 command: str = f"{run_python_cmd_} {root_relative_to_notebooks / file} > {root_relative_to_notebooks / output_file} 2>&1" 169 process: subprocess.CompletedProcess = subprocess.run( 170 command, shell=True, text=True 171 ) 172 173 if process.returncode == 0: 174 print(f" ✅ Run completed with return code {process.returncode}") 175 else: 176 print( 177 f" ❌ Run failed with return code {process.returncode}!!! Check {output_file.as_posix()}" 178 ) 179 180 # print the output of the file to the console if it failed 181 if process.returncode != 0: 182 with open(root_relative_to_notebooks / output_file, "r") as f: 183 file_output: str = f.read() 184 err: str = f"Error in {file}:\n{'-'*term_width}\n{file_output}" 185 exceptions[file.as_posix()] = err 186 if exit_on_first_fail: 187 raise NotebookTestError(err) 188 189 del process 190 191 if len(exceptions) > 0: 192 exceptions_str: str = ("\n" + "=" * term_width + "\n").join( 193 list(exceptions.values()) 194 ) 195 raise NotebookTestError( 196 exceptions_str 197 + "=" * term_width 198 + f"\n❌ {len(exceptions)}/{n_notebooks} notebooks failed:\n{list(exceptions.keys())}" 199 ) 200 201 except NotebookTestError as e: 202 print("!" * term_width, file=sys.stderr) 203 print(e, file=sys.stderr) 204 print("!" * term_width, file=sys.stderr) 205 raise e 206 finally: 207 # return to original cwd 208 os.chdir(original_cwd)
Run converted Jupyter notebooks as Python scripts and verify they execute successfully.
Takes a directory of notebooks and their corresponding converted Python scripts, executes each script, and captures the output. Failures are collected and reported, with optional early exit on first failure.
Parameters:
notebooks_dir : Path
Directory containing the original .ipynb notebook filesconverted_notebooks_temp_dir : Path
Directory containing the corresponding converted .py filesCI_output_suffix : str
Suffix to append to output files capturing execution results (defaults to".CI-output.txt"
)run_python_cmd : str | None
Custom command to run Python scripts. Overrides python_tool and run_python_cmd_fmt if provided (defaults toNone
)run_python_cmd_fmt : str
Format string for constructing the Python run command (defaults to"{python_tool} run python"
)python_tool : str
Tool used to run Python (e.g. poetry, uv) (defaults to"poetry"
)exit_on_first_fail : bool
Whether to raise exception immediately on first notebook failure (defaults toFalse
)
Returns:
None
Modifies:
- Working directory: Temporarily changes to notebooks_dir during execution
- Filesystem: Creates output files with CI_output_suffix for each notebook
Raises:
NotebookTestError
: If any notebooks fail to execute, or if input directories are invalidTypeError
: If run_python_cmd is provided but not a string
Usage:
>>> run_notebook_tests(
... notebooks_dir=Path("notebooks"),
... converted_notebooks_temp_dir=Path("temp/converted"),
... python_tool="poetry"
... )
# testing notebooks in 'notebooks'
# reading converted notebooks from 'temp/converted'
Running 1/2: temp/converted/notebook1.py
Output in temp/converted/notebook1.CI-output.txt
✅ Run completed with return code 0