noxfile.py

#! /usr/bin/env -S uv run -q --script --python 3.11
#
# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2024, ETH Zurich
# All rights reserved.
#
# Please, refer to the LICENSE file in the root directory.
# SPDX-License-Identifier: BSD-3-Clause
#
# Note:
#   The explicit '--python 3.11' in the shebang is only needed due
#   to the existence of the .python-versions file, which overrides
#   the PEP 723 'requires-python' metadata.
# /// script
# requires-python = ">=3.11"
# dependencies = ["nox>=2025.02.09", "uv>=0.6.10"]
# ///

from __future__ import annotations

import pathlib
from collections.abc import Sequence
from typing import Final, Literal, TypeAlias

import nox
import tomllib


# This should just be `pytest.ExitCode.NO_TESTS_COLLECTED` but `pytest`
# is not guaranteed to be available in the venv where `nox` is running.
NO_TESTS_COLLECTED_EXIT_CODE: Final = 5

# -- nox configuration --
nox.options.default_venv_backend = "uv"
nox.options.sessions = [
    "test_cartesian-3.10(internal, cpu)",
    "test_cartesian-3.10(dace, cpu)",
    "test_cartesian-3.11(internal, cpu)",
    "test_cartesian-3.11(dace, cpu)",
    "test_cartesian-3.12(internal, cpu)",
    "test_cartesian-3.12(dace, cpu)",
    "test_cartesian-3.13(internal, cpu)",
    "test_cartesian-3.13(dace, cpu)",
    "test_cartesian-3.14(internal, cpu)",
    "test_cartesian-3.14(dace, cpu)",
    "test_eve-3.10",
    "test_eve-3.11",
    "test_eve-3.12",
    "test_eve-3.13",
    "test_eve-3.14",
    "test_next-3.10(internal, cpu, nomesh)",
    "test_next-3.10(dace, cpu, nomesh)",
    "test_next-3.11(internal, cpu, nomesh)",
    "test_next-3.11(dace, cpu, nomesh)",
    "test_next-3.12(internal, cpu, nomesh)",
    "test_next-3.12(dace, cpu, nomesh)",
    "test_next-3.13(internal, cpu, nomesh)",
    "test_next-3.13(dace, cpu, nomesh)",
    "test_next-3.14(internal, cpu, nomesh)",
    "test_next-3.14(dace, cpu, nomesh)",
    "test_package-3.10",
    "test_package-3.11",
    "test_package-3.12",
    "test_package-3.13",
    "test_package-3.14",
    "test_storage-3.10(cpu)",
    "test_storage-3.11(cpu)",
    "test_storage-3.12(cpu)",
    "test_storage-3.13(cpu)",
    "test_storage-3.14(cpu)",
]

REPO_ROOT: Final = pathlib.Path(__file__).parent.resolve().absolute()
PYTHON_VERSIONS: Final[list[str]] = [
    v
    for line in (REPO_ROOT / ".python-versions").read_text().splitlines()
    if (v := line.strip()) and not v.startswith("#")
]
REQUIRES_PYTHON = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text())["project"][
    "requires-python"
]

# -- Parameter sets --
DeviceOption: TypeAlias = Literal["cpu", "cuda12", "cuda13", "rocm6", "rocm7"]
DeviceNoxParam: Final[dict[DeviceOption, nox.param]] = {
    device: nox.param(device, id=device, tags=[device]) for device in DeviceOption.__args__
}
DeviceTestSettings: Final[dict[DeviceOption, dict[str, list[str]]]] = {
    "cpu": {"extras": [], "markers": ["not requires_gpu"]},
    **{
        device: {"extras": [device], "markers": ["requires_gpu"]}
        for device in DeviceOption.__args__
        if device != "cpu"
    },
}

CodeGenOption: TypeAlias = Literal["internal", "dace"]
CodeGenNoxParam: Final[dict[CodeGenOption, nox.param]] = {
    codegen: nox.param(codegen, id=codegen, tags=[codegen]) for codegen in CodeGenOption.__args__
}
CodeGenTestSettings: Final[dict[str, dict[str, list[str]]]] = {
    "internal": {"extras": ["jax"], "markers": ["not requires_dace"]}
}
# Use dace-cartesian group to select the appropriate dace version
CodeGenCartesianTestSettings = CodeGenTestSettings | {
    "dace": {"extras": [], "groups": ["dace-cartesian"], "markers": ["requires_dace"]},
}
# Install dace-next group to select the appropriate dace version
CodeGenNextTestSettings = CodeGenTestSettings | {
    "dace": {"extras": [], "groups": ["dace-next"], "markers": ["requires_dace"]},
}


# -- Utilities --
def install_session_venv(
    session: nox.Session,
    *args: str | Sequence[str],
    extras: Sequence[str] = (),
    groups: Sequence[str] = (),
) -> None:
    """
    Install session packages using the `uv` tool.

    Args:
        session: The Nox session object.
        *args: Additional packages to install in the session (via `uv pip install`)
        extras: Names of package's extras to install.
        groups: Names of dependency groups to install.
    """
    session.run_install(
        "uv",
        "sync",
        "--python",
        # uv does not yet combine explicit python version requests with the
        # `requires-python` range in `pyproject.toml`, so we do it manually.
        # See: https://github.com/astral-sh/uv/issues/16654
        f"{REQUIRES_PYTHON}, >={session.python!s}.0",
        "--no-dev",
        *(f"--extra={e}" for e in extras),
        *(f"--group={g}" for g in groups),
        env=session.env | dict(UV_PROJECT_ENVIRONMENT=session.virtualenv.location),
    )
    for item in args:
        session.run_install(
            "uv",
            "pip",
            "install",
            *((item,) if isinstance(item, str) else item),
            env=session.env | dict(UV_PROJECT_ENVIRONMENT=session.virtualenv.location),
        )


# -- Sessions --
@nox.session(python=PYTHON_VERSIONS, tags=["cartesian"])
@nox.parametrize("device", [*DeviceNoxParam.values()])
@nox.parametrize("codegen", [*CodeGenNoxParam.values()])
def test_cartesian(
    session: nox.Session,
    codegen: CodeGenOption,
    device: DeviceOption,
) -> None:
    """Run selected 'gt4py.cartesian' tests."""

    codegen_settings = CodeGenCartesianTestSettings[codegen]
    device_settings = DeviceTestSettings[device]
    extras = [
        "standard",
        "testing",
        *codegen_settings.get("extras", []),
        *device_settings.get("extras", []),
    ]
    groups = ["test", *codegen_settings.get("groups", []), *device_settings.get("groups", [])]

    install_session_venv(session, extras=extras, groups=groups)

    markers = " and ".join(codegen_settings["markers"] + device_settings["markers"])

    session.run(
        *"pytest --cache-clear -sv -n auto --dist loadgroup".split(),
        *("-m", f"{markers}"),
        str(pathlib.Path("tests") / "cartesian_tests"),
        *session.posargs,
    )
    session.run(
        *"pytest --doctest-modules --doctest-ignore-import-errors -sv".split(),
        str(pathlib.Path("src") / "gt4py" / "cartesian"),
    )


@nox.session(python=PYTHON_VERSIONS, tags=["cartesian", "next", "cpu"])
def test_eve(session: nox.Session) -> None:
    """Run 'gt4py.eve' tests."""

    install_session_venv(session, groups=["test"])

    session.run(
        *"pytest --cache-clear -sv -n auto --dist loadgroup".split(),
        str(pathlib.Path("tests") / "eve_tests"),
        *session.posargs,
    )
    session.run(
        *"pytest --doctest-modules -sv".split(),
        str(pathlib.Path("src") / "gt4py" / "eve"),
    )


@nox.session(python=PYTHON_VERSIONS, tags=["next"])
def test_examples(session: nox.Session) -> None:
    """Run and test documentation workflows."""

    install_session_venv(session, extras=["testing"], groups=["docs", "test"])

    session.run(*"jupytext docs/user/next/QuickstartGuide.md --to .ipynb".split())
    session.run(*"jupytext docs/user/next/advanced/*.md --to .ipynb".split())

    for notebook, extra_args in [
        ("docs/user/next/workshop/slides", None),
        ("docs/user/next/workshop/exercises", ["-k", "solutions"]),
        ("docs/user/next/QuickstartGuide.ipynb", None),
        ("docs/user/next/advanced", None),
        ("examples", (None)),
    ]:
        session.run(
            *f"pytest --nbmake {notebook} -sv -n 1 --benchmark-disable".split(),
            *(extra_args or []),
        )


@nox.session(python=PYTHON_VERSIONS, tags=["next"])
@nox.parametrize(
    "meshlib",
    [
        nox.param("nomesh", id="nomesh", tags=["nomesh"]),
        nox.param("atlas", id="atlas", tags=["atlas"]),
    ],
)
@nox.parametrize("device", [*DeviceNoxParam.values()])
@nox.parametrize("codegen", [*CodeGenNoxParam.values()])
def test_next(
    session: nox.Session,
    codegen: CodeGenOption,
    device: DeviceOption,
    meshlib: Literal["nomesh", "atlas"],
) -> None:
    """Run selected 'gt4py.next' tests."""

    codegen_settings = CodeGenNextTestSettings[codegen]
    device_settings = DeviceTestSettings[device]
    extras = [
        "standard",
        "testing",
        *codegen_settings.get("extras", []),
        *device_settings.get("extras", []),
    ]
    groups = ["test", *codegen_settings.get("groups", []), *device_settings.get("groups", [])]
    mesh_markers: list[str] = []

    match meshlib:
        case "nomesh":
            mesh_markers.append("not requires_atlas")
        case "atlas":
            mesh_markers.append("requires_atlas")
            groups.append("frameworks")

    install_session_venv(session, extras=extras, groups=groups)

    markers = " and ".join(codegen_settings["markers"] + device_settings["markers"] + mesh_markers)

    session.run(
        *"pytest --cache-clear -sv -n auto --dist loadgroup".split(),
        *("-m", f"{markers}"),
        str(pathlib.Path("tests") / "next_tests"),
        *session.posargs,
        success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE],
    )
    session.run(
        *"pytest --doctest-modules --doctest-ignore-import-errors -sv".split(),
        str(pathlib.Path("src") / "gt4py" / "next"),
        success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE],
    )


@nox.session(python=PYTHON_VERSIONS, tags=["cartesian", "next", "cpu"])
def test_package(session: nox.Session) -> None:
    """Run 'gt4py' package level tests."""

    install_session_venv(session, groups=["test"])

    session.run(
        *"pytest --cache-clear -sv".split(),
        str(pathlib.Path("tests") / "package_tests"),
        *session.posargs,
    )

    modules = [str(path) for path in (pathlib.Path("src") / "gt4py").glob("*.py")]
    session.run(
        *"pytest --doctest-modules --doctest-ignore-import-errors -sv".split(),
        *modules,
        success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE],
    )


@nox.session(python=PYTHON_VERSIONS, tags=["cartesian", "next"])
@nox.parametrize("device", [*DeviceNoxParam.values()])
def test_storage(
    session: nox.Session,
    device: DeviceOption,
) -> None:
    """Run selected 'gt4py.storage' tests."""

    device_settings = DeviceTestSettings[device]

    install_session_venv(
        session, extras=["standard", "testing", *device_settings["extras"]], groups=["test"]
    )

    markers = " and ".join(device_settings["markers"])

    session.run(
        *"pytest --cache-clear -sv -n auto --dist loadgroup".split(),
        *("-m", f"{markers}"),
        str(pathlib.Path("tests") / "storage_tests"),
        *session.posargs,
    )
    session.run(
        *"pytest --doctest-modules -sv".split(),
        str(pathlib.Path("src") / "gt4py" / "storage"),
        success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE],
    )


@nox.session(python=PYTHON_VERSIONS, tags=["next"])
def test_typing_exports(session: nox.Session) -> None:
    """Test GT4Py usability in a typed client context."""
    install_session_venv(session, extras=["standard"], groups=["test", "typing_exports"])

    session.run(
        "pytest",
        "-sv",
        "--mypy-testing-base",
        "typing_tests",
        "typing_tests",
        *session.posargs,
    )


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