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