Skip to content

Commit

Permalink
print out the successes as well
Browse files Browse the repository at this point in the history
  • Loading branch information
wd60622 committed Nov 16, 2024
1 parent cdfb6ab commit 059d628
Showing 1 changed file with 23 additions and 18 deletions.
41 changes: 23 additions & 18 deletions scripts/run_notebooks/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
from argparse import ArgumentParser

from rich.console import Console
from dataclasses import dataclass
import logging
from pathlib import Path
from tempfile import NamedTemporaryFile
Expand Down Expand Up @@ -99,7 +100,13 @@ def actual_run(notebook_path: Path, i: int, total: int) -> None:
)


class NotebookFailure(TypedDict):
@dataclass
class NotebookSuccess:
notebook_path: Path


@dataclass
class NotebookFailure:
notebook_path: Path
error: str

Expand All @@ -109,19 +116,17 @@ def run_notebook(
i: int,
total: int,
mock: bool = True,
) -> NotebookFailure | None:
) -> NotebookFailure | NotebookSuccess:
logging.info(f"Running notebook: {notebook_path.name}")
run = mock_run if mock else actual_run

try:
run(notebook_path, i=i, total=total)
except Exception as e:
logging.error(
f"{e.__class__.__name__} encountered running notebook: {str(notebook_path)}"
)
logging.error(f"{e.__class__.__name__} encountered running notebook: {str(notebook_path)}")
return NotebookFailure(notebook_path=notebook_path, error=str(e))
else:
return
return NotebookSuccess(notebook_path=notebook_path)


class RunParams(TypedDict):
Expand All @@ -133,35 +138,35 @@ class RunParams(TypedDict):

def run_parameters(notebook_paths: list[Path], mock: bool = True) -> list[RunParams]:
def to_mock(notebook_path: Path, i: int) -> RunParams:
return RunParams(
notebook_path=notebook_path, mock=mock, i=i, total=len(notebook_paths)
)
return RunParams(notebook_path=notebook_path, mock=mock, i=i, total=len(notebook_paths))

return [
to_mock(notebook_path, i=i)
for i, notebook_path in enumerate(notebook_paths, start=1)
]
return [to_mock(notebook_path, i=i) for i, notebook_path in enumerate(notebook_paths, start=1)]


def main(notebooks_to_run: list[Path], mock: bool = True) -> None:
console = Console()
errors: list[NotebookFailure]
setup_logging()
logging.info("Starting notebook runner")
logging.info(f"Running {len(notebooks_to_run)} notebook(s).")
results = Parallel(n_jobs=-1)(
delayed(run_notebook)(**run_params)
for run_params in run_parameters(notebooks_to_run, mock=mock)
)
errors = [result for result in results if result is not None]
errors: list[NotebookFailure] = list(filter(lambda x: isinstance(x, NotebookFailure), results))
successes: list[NotebookSuccess] = list(
filter(lambda x: isinstance(x, NotebookSuccess), results)
)

if not errors:
logging.info("Notebooks run successfully!")
logging.info("All notebooks ran successfully!")
return

for error in errors:
console.rule(f"[bold red]Error running {error['notebook_path']}[/bold red]")
console.print(error["error"])
console.rule(f"[bold red]Error running {error.notebook_path}[/bold red]")
console.print(error.error)

for success in successes:
console.print(f"[bold green]Success running {success.notebook_path}[/bold green]")

logging.error(f"{len(errors)} / {len(notebooks_to_run)} notebooks failed")

Expand Down

0 comments on commit 059d628

Please sign in to comment.