docs for muutils v0.6.20
View Source on GitHub

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
 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    )

class NotebookTestError(builtins.Exception):
21class NotebookTestError(Exception):
22    pass

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: Optional[str] = None, run_python_cmd_fmt: str = '{python_tool} run python', python_tool: str = 'poetry', exit_on_first_fail: bool = False):
 25def run_notebook_tests(
 26    notebooks_dir: Path,
 27    converted_notebooks_temp_dir: Path,
 28    CI_output_suffix: str = ".CI-output.txt",
 29    run_python_cmd: Optional[str] = None,
 30    run_python_cmd_fmt: str = "{python_tool} run python",
 31    python_tool: str = "poetry",
 32    exit_on_first_fail: bool = False,
 33):
 34    """Run converted Jupyter notebooks as Python scripts and verify they execute successfully.
 35
 36    Takes a directory of notebooks and their corresponding converted Python scripts,
 37    executes each script, and captures the output. Failures are collected and reported,
 38    with optional early exit on first failure.
 39
 40    # Parameters:
 41     - `notebooks_dir : Path`
 42        Directory containing the original .ipynb notebook files
 43     - `converted_notebooks_temp_dir : Path`
 44        Directory containing the corresponding converted .py files
 45     - `CI_output_suffix : str`
 46        Suffix to append to output files capturing execution results
 47        (defaults to `".CI-output.txt"`)
 48     - `run_python_cmd : str | None`
 49        Custom command to run Python scripts. Overrides python_tool and run_python_cmd_fmt if provided
 50        (defaults to `None`)
 51     - `run_python_cmd_fmt : str`
 52        Format string for constructing the Python run command
 53        (defaults to `"{python_tool} run python"`)
 54     - `python_tool : str`
 55        Tool used to run Python (e.g. poetry, uv)
 56        (defaults to `"poetry"`)
 57     - `exit_on_first_fail : bool`
 58        Whether to raise exception immediately on first notebook failure
 59        (defaults to `False`)
 60
 61    # Returns:
 62     - `None`
 63
 64    # Modifies:
 65     - Working directory: Temporarily changes to notebooks_dir during execution
 66     - Filesystem: Creates output files with CI_output_suffix for each notebook
 67
 68    # Raises:
 69     - `NotebookTestError`: If any notebooks fail to execute, or if input directories are invalid
 70     - `TypeError`: If run_python_cmd is provided but not a string
 71
 72    # Usage:
 73    ```python
 74    >>> run_notebook_tests(
 75    ...     notebooks_dir=Path("notebooks"),
 76    ...     converted_notebooks_temp_dir=Path("temp/converted"),
 77    ...     python_tool="poetry"
 78    ... )
 79    # testing notebooks in 'notebooks'
 80    # reading converted notebooks from 'temp/converted'
 81    Running 1/2: temp/converted/notebook1.py
 82        Output in temp/converted/notebook1.CI-output.txt
 83        ✅ Run completed with return code 0
 84    ```
 85    """
 86
 87    run_python_cmd_: str
 88    if run_python_cmd is None:
 89        run_python_cmd_ = run_python_cmd_fmt.format(python_tool=python_tool)
 90    elif isinstance(run_python_cmd, str):
 91        run_python_cmd_ = run_python_cmd
 92        warnings.warn(
 93            "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.",
 94            DeprecationWarning,
 95        )
 96    else:
 97        raise TypeError(
 98            f"run_python_cmd must be a string or None, got {run_python_cmd =}, {type(run_python_cmd) =}"
 99        )
100
101    original_cwd: Path = Path.cwd()
102    # get paths
103    notebooks_dir = Path(notebooks_dir)
104    converted_notebooks_temp_dir = Path(converted_notebooks_temp_dir)
105    root_relative_to_notebooks: Path = Path(os.path.relpath(".", notebooks_dir))
106
107    term_width: int
108    try:
109        term_width = os.get_terminal_size().columns
110    except OSError:
111        term_width = 80
112
113    exceptions: dict[str, str] = dict()
114
115    print(f"# testing notebooks in '{notebooks_dir}'")
116    print(
117        f"# reading converted notebooks from '{converted_notebooks_temp_dir.as_posix()}'"
118    )
119
120    try:
121        # check things exist
122        if not notebooks_dir.exists():
123            raise NotebookTestError(f"Notebooks dir '{notebooks_dir}' does not exist")
124        if not notebooks_dir.is_dir():
125            raise NotebookTestError(
126                f"Notebooks dir '{notebooks_dir}' is not a directory"
127            )
128        if not converted_notebooks_temp_dir.exists():
129            raise NotebookTestError(
130                f"Converted notebooks dir '{converted_notebooks_temp_dir}' does not exist"
131            )
132        if not converted_notebooks_temp_dir.is_dir():
133            raise NotebookTestError(
134                f"Converted notebooks dir '{converted_notebooks_temp_dir}' is not a directory"
135            )
136
137        notebooks: list[Path] = list(notebooks_dir.glob("*.ipynb"))
138        if not notebooks:
139            raise NotebookTestError(f"No notebooks found in '{notebooks_dir}'")
140
141        converted_notebooks: list[Path] = list()
142        for nb in notebooks:
143            converted_file: Path = (
144                converted_notebooks_temp_dir / nb.with_suffix(".py").name
145            )
146            if not converted_file.exists():
147                raise NotebookTestError(
148                    f"Did not find converted notebook '{converted_file}' for '{nb}'"
149                )
150            converted_notebooks.append(converted_file)
151
152        del converted_file
153
154        # the location of this line is important
155        os.chdir(notebooks_dir)
156
157        n_notebooks: int = len(converted_notebooks)
158        for idx, file in enumerate(converted_notebooks):
159            # run the file
160            print(f"Running {idx+1}/{n_notebooks}: {file.as_posix()}")
161            output_file: Path = file.with_suffix(CI_output_suffix)
162            print(f"    Output in {output_file.as_posix()}")
163            with SpinnerContext(
164                spinner_chars="braille",
165                update_interval=0.5,
166                format_string="\r    {spinner} ({elapsed_time:.2f}s) {message}{value}",
167            ):
168                command: str = f"{run_python_cmd_} {root_relative_to_notebooks / file} > {root_relative_to_notebooks / output_file} 2>&1"
169                process: subprocess.CompletedProcess = subprocess.run(
170                    command, shell=True, text=True
171                )
172
173            if process.returncode == 0:
174                print(f"    ✅ Run completed with return code {process.returncode}")
175            else:
176                print(
177                    f"    ❌ Run failed with return code {process.returncode}!!! Check {output_file.as_posix()}"
178                )
179
180            # print the output of the file to the console if it failed
181            if process.returncode != 0:
182                with open(root_relative_to_notebooks / output_file, "r") as f:
183                    file_output: str = f.read()
184                err: str = f"Error in {file}:\n{'-'*term_width}\n{file_output}"
185                exceptions[file.as_posix()] = err
186                if exit_on_first_fail:
187                    raise NotebookTestError(err)
188
189            del process
190
191        if len(exceptions) > 0:
192            exceptions_str: str = ("\n" + "=" * term_width + "\n").join(
193                list(exceptions.values())
194            )
195            raise NotebookTestError(
196                exceptions_str
197                + "=" * term_width
198                + f"\n{len(exceptions)}/{n_notebooks} notebooks failed:\n{list(exceptions.keys())}"
199            )
200
201    except NotebookTestError as e:
202        print("!" * term_width, file=sys.stderr)
203        print(e, file=sys.stderr)
204        print("!" * term_width, file=sys.stderr)
205        raise e
206    finally:
207        # return to original cwd
208        os.chdir(original_cwd)

Run converted Jupyter notebooks as Python scripts and verify they execute successfully.

Takes a directory of notebooks and their corresponding converted Python scripts, executes each script, and captures the output. Failures are collected and reported, with optional early exit on first failure.

Parameters:

  • notebooks_dir : Path Directory containing the original .ipynb notebook files
  • converted_notebooks_temp_dir : Path Directory containing the corresponding converted .py files
  • CI_output_suffix : str Suffix to append to output files capturing execution results (defaults to ".CI-output.txt")
  • run_python_cmd : str | None Custom command to run Python scripts. Overrides python_tool and run_python_cmd_fmt if provided (defaults to None)
  • run_python_cmd_fmt : str Format string for constructing the Python run command (defaults to "{python_tool} run python")
  • python_tool : str Tool used to run Python (e.g. poetry, uv) (defaults to "poetry")
  • exit_on_first_fail : bool Whether to raise exception immediately on first notebook failure (defaults to False)

Returns:

  • None

Modifies:

  • Working directory: Temporarily changes to notebooks_dir during execution
  • Filesystem: Creates output files with CI_output_suffix for each notebook

Raises:

  • NotebookTestError: If any notebooks fail to execute, or if input directories are invalid
  • TypeError: If run_python_cmd is provided but not a string

Usage:

>>> run_notebook_tests(
...     notebooks_dir=Path("notebooks"),
...     converted_notebooks_temp_dir=Path("temp/converted"),
...     python_tool="poetry"
... )
# testing notebooks in 'notebooks'
# reading converted notebooks from 'temp/converted'
Running 1/2: temp/converted/notebook1.py
    Output in temp/converted/notebook1.CI-output.txt
     Run completed with return code 0