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