diff --git a/pyproject.toml b/pyproject.toml index 1dc8e9e..6111b25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,24 @@ features = ["test"] [tool.hatch.envs.test.scripts] run = "pytest -vv {args}" +[tool.ruff] +select = ["ALL"] +extend-ignore = [ + "COM812", "ISC001", # Incompatible with formatter + "D203", # Use no blank line before class + "D213", # Use multiline summary on first line + "FIX002", # Comments with “TODO” are OK + "TD002", # No need for TODO authors +] +allowed-confusables = ["’", "×"] +[tool.ruff.extend-per-file-ignores] +"tests/**/*.py" = [ + "D100", # No module docstrings necessary for tests + "D103", # No function docstrings necessary for tests + "INP001", # Pytest directories are not packages + "S101", # Pytest rewrites asserts +] + [build-system] requires = ["hatchling", "hatch-vcs"] build-backend = "hatchling.build" diff --git a/src/prelude_runner/__init__.py b/src/prelude_runner/__init__.py index 6bbc0e3..1079145 100644 --- a/src/prelude_runner/__init__.py +++ b/src/prelude_runner/__init__.py @@ -1,3 +1,4 @@ +"""Prelude Runner.""" from .core import Preludes, execute __all__ = ["execute", "Preludes"] diff --git a/src/prelude_runner/cli.py b/src/prelude_runner/cli.py index 953e5b7..291fedb 100644 --- a/src/prelude_runner/cli.py +++ b/src/prelude_runner/cli.py @@ -1,14 +1,20 @@ +"""Command line parsing and execution.""" +from __future__ import annotations + from argparse import ArgumentParser from pathlib import Path -from typing import Protocol +from typing import TYPE_CHECKING, Protocol, Sequence import nbformat from .core import Preludes, execute -from .types import Notebook + +if TYPE_CHECKING: + from .types import Notebook def load_preludes(d: Path) -> Preludes: + """Load prelude code from config directory.""" return Preludes( notebook=(d / "prelude_notebook.py").read_text(), cell=(d / "prelude_cell.py").read_text(), @@ -16,18 +22,22 @@ def load_preludes(d: Path) -> Preludes: class Args(Protocol): + """Parsed command line arguments.""" + preludes: Path nb_path: Path -def parse_args(argv: list[str] | None = None) -> Args: - parser = ArgumentParser(argv) +def parse_args(argv: Sequence[str] | None = None) -> Args: + """Parse command line arguments.""" + parser = ArgumentParser() parser.add_argument("--preludes", type=Path, help="Path to prelude directory") parser.add_argument("nb-path", type=Path, help="Path to notebook directory") - return parser.parse_args() + return parser.parse_args(argv) -def main(argv: list[str] | None = None) -> None: +def main(argv: Sequence[str] | None = None) -> None: + """Execute main entry point.""" args = parse_args(argv) preludes = load_preludes(args.preludes) for nb_path in args.nb_path.rglob("*.ipynb"): diff --git a/src/prelude_runner/core.py b/src/prelude_runner/core.py index 292c240..65a2a49 100644 --- a/src/prelude_runner/core.py +++ b/src/prelude_runner/core.py @@ -1,32 +1,45 @@ +"""Main exports.""" +from __future__ import annotations + from dataclasses import dataclass -from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any -from jupyter_client.manager import KernelManager from nbclient import NotebookClient from nbclient.util import ensure_async -from .types import CodeCell, Notebook +if TYPE_CHECKING: + from pathlib import Path + + from jupyter_client.manager import KernelManager + + from .types import CodeCell, Notebook @dataclass class Preludes: + """Prelude code to execute before notebook and cell.""" + notebook: str | None = None cell: str | None = None -def add_reproducible_hooks(client: NotebookClient, preludes: Preludes) -> None: +def add_prelude_hooks(client: NotebookClient, preludes: Preludes) -> None: + """Add hooks for prelude execution to NotebookClient.""" + async def execute(source: str) -> None: await ensure_async( client.kc.execute( - source, silent=True, store_history=False, allow_stdin=False - ) + source, + silent=True, + store_history=False, + allow_stdin=False, + ), ) - async def on_notebook_start(notebook: Notebook) -> None: + async def on_notebook_start(notebook: Notebook) -> None: # noqa: ARG001 await execute(preludes.notebook) - async def on_cell_execute(cell: CodeCell, cell_index: int) -> None: + async def on_cell_execute(cell: CodeCell, cell_index: int) -> None: # noqa: ARG001 await execute(preludes.cell) client.on_notebook_start = on_notebook_start @@ -39,12 +52,12 @@ def execute( *, cwd: Path | None = None, km: KernelManager | None = None, - **kwargs: Any, + **kwargs: Any, # noqa: ANN401 ) -> Notebook: """Execute a notebook's code, updating outputs within the notebook object.""" resources = {} if cwd is not None: resources["metadata"] = {"path": cwd} client = NotebookClient(nb=nb, resources=resources, km=km, **kwargs) - add_reproducible_hooks(client, preludes) + add_prelude_hooks(client, preludes) return client.execute() diff --git a/src/prelude_runner/types.py b/src/prelude_runner/types.py index b7e613b..36207a5 100644 --- a/src/prelude_runner/types.py +++ b/src/prelude_runner/types.py @@ -1,38 +1,60 @@ +"""Types more specific than NotebookNode.""" +from __future__ import annotations + from typing import Any, Literal, Protocol, TypedDict class CellMetadata(TypedDict): + """Common metadata for all cells.""" + tags: list[str] class Cell(Protocol): + """A cell in a notebook.""" + cell_type: Literal["markdown", "code"] metadata: CellMetadata source: str class Output(Protocol): + """A code cell’s output.""" + output_type: Literal[ - "stream", "display_data", "execute_result", "error", "update_display_data" + "stream", + "display_data", + "execute_result", + "error", + "update_display_data", ] data: dict[str, str] class Stream(Protocol): + """One type of output from a code cell.""" + output_type: Literal["stream"] name: Literal["stdout", "stderr"] text: str class CodeCell(Cell): + """A code cell.""" + cell_type: Literal["code"] - outputs: list[Stream | Any] # TODO + # TODO: add other output types # noqa: TD003 + outputs: list[Stream | Any] class NotebookMetadata(TypedDict): + """Common metadata for all notebooks.""" + language_info: dict[str, str] class Notebook(Protocol): + """A Jupyter notebook.""" + metadata: NotebookMetadata cells: list[Cell] diff --git a/tests/data/config/prelude_cell.py b/tests/data/config/prelude_cell.py new file mode 100644 index 0000000..104297a --- /dev/null +++ b/tests/data/config/prelude_cell.py @@ -0,0 +1,9 @@ +import random +from contextlib import suppress + +random.seed(0) + +with suppress(Exception): + from numpy.random import seed + + seed(0) # noqa: NPY002 diff --git a/tests/example-config/prelude_notebook.py b/tests/data/config/prelude_notebook.py similarity index 100% rename from tests/example-config/prelude_notebook.py rename to tests/data/config/prelude_notebook.py diff --git a/tests/data/notebooks/example.ipynb b/tests/data/notebooks/example.ipynb new file mode 100644 index 0000000..e69de29 diff --git a/tests/example-config/prelude_cell.py b/tests/example-config/prelude_cell.py deleted file mode 100644 index 7832b40..0000000 --- a/tests/example-config/prelude_cell.py +++ /dev/null @@ -1,16 +0,0 @@ -import random -from contextlib import suppress - -random.seed(0) - -with suppress(Exception): - from numpy.random import seed - - seed(0) - -with suppress(Exception): - # https://pytorch.org/docs/stable/notes/randomness.html - import torch - - torch.use_deterministic_algorithms(True) - torch.manual_seed(0) diff --git a/tests/test_runner.py b/tests/test_runner.py index 0ed1608..c27c0a1 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,19 +1,22 @@ from pathlib import Path +from typing import TYPE_CHECKING import nbclient import pytest from nbclient.exceptions import CellExecutionError from nbformat import v4 -from prelude_runner.cli import load_preludes +from prelude_runner.cli import load_preludes, main from prelude_runner.core import Preludes, execute -from prelude_runner.types import CodeCell, Notebook + +if TYPE_CHECKING: + from prelude_runner.types import CodeCell, Notebook tests_dir = Path(__file__).parent @pytest.fixture(scope="session") -def preludes(): - return load_preludes(tests_dir / "example-config") +def preludes() -> Preludes: + return load_preludes(tests_dir / "data/config") @pytest.mark.parametrize( @@ -21,21 +24,23 @@ def preludes(): [ pytest.param("import random; print(random.randint(0, 10))", "6", id="stdlib"), pytest.param( - "import numpy as np; print(np.random.randint(0, 10))", "5", id="numpy" + "import numpy as np; print(np.random.randint(0, 10))", + "5", + id="numpy", ), ], ) -def test_execute(preludes: Preludes, code: str, expected: str): +def test_execute(preludes: Preludes, code: str, expected: str) -> None: cell: CodeCell = v4.new_code_cell(cell_type="code", source=code) nb: Notebook = v4.new_notebook(cells=[cell]) execute(nb, preludes) assert cell.outputs[0].text.strip() == expected -def test_traceback_intact(preludes: Preludes): +def test_traceback_intact(preludes: Preludes) -> None: """Tests that the traceback reports the same line and cell numbers.""" - def mk_nb(): + def mk_nb() -> None: cell: CodeCell = v4.new_code_cell(cell_type="code", source="1/0") nb: Notebook = v4.new_notebook(cells=[cell]) return nb @@ -45,3 +50,8 @@ def mk_nb(): with pytest.raises(CellExecutionError) as exc_orig: nbclient.execute(mk_nb()) assert exc_rr.value.traceback == exc_orig.value.traceback + + +def test_cli() -> None: + nb_path = tests_dir / "data/notebooks" + main([nb_path])