1"""turn a folder of notebooks into scripts, run them, and make sure they work.

2

3 made to be called as

4

5```bash

6 python -m muutils.nbutils.run_notebook_tests --notebooks-dir <notebooks_dir> --converted-notebooks-temp-dir <converted_notebooks_temp_dir>

7```

8"""

9

10 import os

11 import subprocess

12 import sys

13 from pathlib import Path

14

15 from muutils.spinner import SpinnerContext

16

17

18 class NotebookTestError(Exception):

19 pass

20

21

22 def 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

139 if __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 )