Coverage for muutils/nbutils/run_notebook_tests.py: 80%

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 )