Coverage for muutils\nbutils\run_notebook_tests.py: 65%

74 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-15 20:56 -0600

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 )