noxfile.py

import os
import sys
from pathlib import Path

import nox

# Options to modify nox behaviour
nox.options.default_venv_backend = "uv|virtualenv"
nox.options.reuse_existing_virtualenvs = True
nox.needs_version = ">= 2025.10.14"

homedir = os.getenv("HOME")
PYBAMM_ENV = {
    "PYTHONIOENCODING": "utf-8",
    "MPLBACKEND": "Agg",
}
VENV_DIR = Path("./venv").resolve()


def set_environment_variables(env_dict, session):
    """
    Sets environment variables for a nox Session object.

    Parameters
    -----------
        session : nox.Session
            The session to set the environment variables for.
        env_dict : dict
            A dictionary of environment variable names and values.

    """
    for key, value in env_dict.items():
        session.env[key] = value


def install_locked(session, *, extras=None, groups=None):
    """Install pybamm and its dependencies into the session environment.

    Two modes, selected by the ``PYBAMM_SOLVER_WHEELS`` environment variable:

    * **Unset (local dev, solver CI):** ``uv sync --frozen`` over the workspace.
      ``pybammsolvers`` is built from the in-repo source via the shared lockfile.

    * **Set to a directory of prebuilt solver wheels (the PyBaMM CI matrix):**
      install the wheel matching this interpreter by explicit path, then install
      pybamm with ``--no-sources`` so the workspace source routing is ignored.
      This tests PyBaMM against the *prebuilt in-repo* ``pybammsolvers`` without
      recompiling it in every matrix cell (and without needing a from-source
      build on Windows, which only exists via cibuildwheel/vcpkg).

      ``--no-sources`` plus an explicit wheel path is required because the in-repo
      solver version collides with the PyPI release, so neither ``--find-links``
      resolution nor the workspace source can be relied on to select the in-repo
      artifact over the identically-versioned PyPI wheel.
    """
    env = {"UV_PROJECT_ENVIRONMENT": session.virtualenv.location}

    wheels = os.getenv("PYBAMM_SOLVER_WHEELS")
    if wheels:
        wheels_path = Path(wheels)
        if wheels_path.is_dir():
            # The per-OS artifact holds wheels for every Python version; pick the
            # one whose interpreter tag matches this cell. The trailing dash in
            # "*-cpXY-*" avoids matching free-threaded "cpXYt-" builds.
            tag = f"cp{sys.version_info.major}{sys.version_info.minor}"
            matches = sorted(wheels_path.glob(f"*-{tag}-*.whl"))
            if not matches:
                session.error(
                    f"PYBAMM_SOLVER_WHEELS={wheels} contains no pybammsolvers "
                    f"wheel for {tag}"
                )
            wheel = matches[0].as_posix()
        else:
            wheel = wheels_path.as_posix()
        python_bin = os.path.join(
            session.bin, "python.exe" if sys.platform == "win32" else "python"
        )
        extras_str = f"[{','.join(extras)}]" if extras else ""
        cmd = [
            "uv",
            "pip",
            "install",
            "--python",
            python_bin,
            "--no-sources",
            wheel,
            "-e",
            f"./packages/pybamm{extras_str}",
        ]
        for group in groups or []:
            # Groups (dev, docs) are defined in the pybamm package, not the root.
            cmd.extend(["--group", f"packages/pybamm/pyproject.toml:{group}"])
        session.run(*cmd, env=env, external=True)
        return

    cmd = ["uv", "sync", "--frozen"]
    for extra in extras or []:
        cmd.extend(["--extra", extra])
    for group in groups or []:
        cmd.extend(["--group", group])
    session.run(*cmd, env=env, external=True)


@nox.session(name="coverage", default=False)
def run_coverage(session):
    """Run the coverage tests and generate an XML report."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax"], groups=["dev"])
    # Using plugin here since coverage runs unit tests on linux with latest python version.
    if "CI" in os.environ:
        session.install("pytest-github-actions-annotate-failures")
    session.run(
        "pytest", "--cov=pybamm", "--cov-report=xml", "packages/pybamm/tests/unit"
    )


@nox.session(name="integration", default=False)
def run_integration(session):
    """Run the integration tests."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax", "pydiffsol"], groups=["dev"])
    if (
        "CI" in os.environ
        and sys.version_info[:2] >= (3, 12)
        and sys.platform == "linux"
    ):
        session.install("pytest-github-actions-annotate-failures")
    session.run("python", "-m", "pytest", "-m", "integration", "packages/pybamm/tests")


@nox.session(name="doctests", default=False)
def run_doctests(session):
    """Run the doctests and generate the output(s) in the docs/build/ directory."""
    install_locked(session, extras=["all"], groups=["dev", "docs"])
    # Fix for Python 3.12 CI. This can be removed after pybtex is replaced.
    session.install("setuptools", silent=False)
    session.run(
        "python",
        "-m",
        "pytest",
        "--doctest-plus",
        "packages/pybamm/src",
    )


@nox.session(name="unit", default=True)
def run_unit(session):
    """Run the unit tests."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax"], groups=["dev"])
    session.run("python", "-m", "pytest", "-m", "unit", "packages/pybamm/tests")


@nox.session(name="memory", default=False)
def run_memory(session):
    """Run memory leak tests using pytest-memray (Linux/macOS only)."""
    if sys.platform == "win32":
        session.skip("memray is not supported on Windows")
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, groups=["dev"])
    session.run(
        "python",
        "-m",
        "pytest",
        "packages/pybamm/tests/memory/",
        "-v",
        "-o",
        "addopts=",
    )


@nox.session(name="examples", default=False)
def run_examples(session):
    """Run the examples tests for Jupyter notebooks."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax"], groups=["dev"])
    notebooks_to_test = session.posargs if session.posargs else []
    session.run(
        "pytest", "--nbmake", *notebooks_to_test, "docs/source/examples/", external=True
    )


@nox.session(name="scripts", default=False)
def run_scripts(session):
    """Run the scripts tests for Python scripts."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax"], groups=["dev"])
    # Fix for Python 3.12 CI. This can be removed after pybtex is replaced.
    session.install("setuptools", silent=False)
    session.run("python", "-m", "pytest", "-m", "scripts", "packages/pybamm/tests")


@nox.session(name="dev", default=False)
def set_dev(session):
    """Install PyBaMM in editable mode."""
    set_environment_variables(PYBAMM_ENV, session=session)
    session.run(
        "uv",
        "sync",
        "--frozen",
        "--extra",
        "all",
        "--extra",
        "jax",
        "--group",
        "dev",
        env={"UV_PROJECT_ENVIRONMENT": os.fsdecode(VENV_DIR)},
        external=True,
    )


@nox.session(name="tests", default=False)
def run_tests(session):
    """Run the unit tests and integration tests sequentially."""
    set_environment_variables(PYBAMM_ENV, session=session)
    install_locked(session, extras=["all", "jax"], groups=["dev"])
    session.run(
        "python",
        "-m",
        "pytest",
        *(
            session.posargs
            if session.posargs
            else ["-m", "unit or integration", "packages/pybamm/tests"]
        ),
    )


@nox.session(name="docs", default=False)
def build_docs(session):
    """Build the documentation and load it in a browser tab, rebuilding on changes."""
    envbindir = session.bin
    install_locked(session, extras=["all"], groups=["docs"])
    # Fix for Python 3.12 CI. This can be removed after pybtex is replaced.
    session.install("setuptools", silent=False)
    session.chdir("docs")
    # Local development
    if session.interactive:
        session.run(
            "sphinx-autobuild",
            "-j",
            "auto",
            "--open-browser",
            "-qT",
            ".",
            f"{envbindir}/../tmp/html",
        )
    # Runs in CI only, treating warnings as errors
    # Run in single-threaded mode, see
    # https://github.com/pydata/pydata-sphinx-theme/issues/1643
    else:
        session.run(
            "sphinx-build",
            "-b",
            "html",
            "-W",
            "--keep-going",
            ".",
            f"{envbindir}/../tmp/html",
        )


@nox.session(name="pre-commit", default=True)
def lint(session):
    """Check all files against the defined pre-commit hooks."""
    session.install("pre-commit", silent=False)
    session.run("pre-commit", "run", "--all-files")


@nox.session(name="quick", reuse_venv=True, default=False)
def run_quick(session):
    """Run integration tests, unit tests, and doctests sequentially"""
    run_tests(session)
    run_doctests(session)


if __name__ == "__main__":
    nox.main()