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

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 )