docs for muutils v0.8.2
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.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    )

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

Common base class for all non-exit exceptions.

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