feat: migrate to nox, run tests in parallel, revamp test infra

.github/workflows/check.yml

@@ -1,4 +1,4 @@
-name: Checks
+name: Style and package checks
 
 on:
   pull_request:
@@ -7,27 +7,31 @@ on:
   push:
     branches:
     - master
+  workflow_dispatch:
 
 env:
-  PIP_DISABLE_PIP_VERSION_CHECK: '1'
-  PY_COLORS: '1'
+  PIP_DISABLE_PIP_VERSION_CHECK: "1"
+  FORCE_COLOR: "3"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
 
 jobs:
   check:
+    name: ${{ matrix.env }}
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        env:
-        #- ruff
-        - package
+        session:
+      # - lint
+        - validate-package
     steps:
     - uses: actions/checkout@v4
     - uses: actions/setup-python@v5.1.1
-      with:
-        python-version: '3.10'
-    - name: Install build tools
-      run: |
-        python -m pip install tox wheel
+
+    - uses: yezz123/setup-uv@v4
+
     - name: Run ${{ matrix.env }}
-      run: python -m tox -e ${{ matrix.env }}
+      run: uvx nox -s ${{ matrix.env }}

.github/workflows/test.yml

@@ -1,53 +1,46 @@
-name: Tests
+name: CI
 
 on:
   pull_request:
     branches:
-    - master
+      - master
   push:
     branches:
-    - master
+      - master
   workflow_dispatch:
   schedule:
-    - cron: '0 4 * * *'
+    - cron: "0 4 * * *"
 
 env:
-  PIP_DISABLE_PIP_VERSION_CHECK: '1'
-  PY_COLORS: '1'
+  PIP_DISABLE_PIP_VERSION_CHECK: "1"
+  FORCE_COLOR: "3"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
 
 jobs:
   test:
+    name: Test / ${{ matrix.platform }} / Nightly ${{ matrix.nightly[0] }} / Python ${{ matrix.python-version }}
     runs-on: ${{ matrix.platform }}
     strategy:
       fail-fast: false
       matrix:
-        platform:
-        - ubuntu-latest
-        - macos-latest
-        - windows-latest
+        platform: [ubuntu-latest, macos-13, macos-latest, windows-latest]
         python-version:
-        - '3.9'
-        - '3.10'
-        - '3.11'
-        - '3.12'
-        - pypy-3.9
-        - pypy-3.10
-        # TODO: bring this back later
-        exclude:
-        - platform: windows-latest
-          python-version: 'pypy-3.10'
+          ["3.8", "3.9", "3.10", "3.11", "3.12", "pypy-3.9", "pypy-3.10"]
+        nightly: [[True, "nightly-"], [False, ""]]
     steps:
-    - uses: actions/checkout@v4
-    - uses: actions/setup-python@v5.1.1
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install Python build tools
-      run: python -m pip install tox wheel
-    - name: Run tests
-      run: python -m tox run -e py
-    - name: Install Scipy prerequisites for Ubuntu
-      if: startsWith(matrix.platform, 'ubuntu')
-      run: sudo apt-get install libopenblas-dev
-    - name: Run tests with scipy
-      if: startsWith(matrix.platform, 'ubuntu') || startsWith(matrix.python-version, 'pypy') != true
-      run: python -m tox run -e py-scipy
+      - uses: actions/checkout@v4.1.7
+      - uses: actions/setup-python@v5.1.1
+        with:
+          python-version: ${{ matrix.python-version }}
+      - uses: yezz123/setup-uv@v4
+
+      - name: Run CPython tests
+        if: ${{ !startsWith(matrix.python-version, 'pypy') }}
+        run: uvx nox -s ${{ matrix.nightly[1] }}tests
+
+      - name: Run PyPy tests
+        if: ${{ startsWith(matrix.python-version, 'pypy') }}
+        run: uvx nox -s ${{ matrix.nightly[1] }}tests

.gitignore

@@ -28,6 +28,7 @@ pip-delete-this-directory.txt
 # Unit test / coverage reports
 htmlcov/
 .tox/
+.nox/
 .coverage
 .coverage.*
 .cache

CONTRIBUTING.md

@@ -1,38 +1,34 @@
 # Contributing
 
-Use Tox to run tests and linting, e.g.
+Use [Nox](https://nox.thea.codes/en/stable/) to run tests and linting, e.g.,
 
 ```shell
-pip install tox
+pip install nox
 ```
 
+`nox` will run all checks in an isolated virtual environment with Autograd and its dependencies, including its optional dependencies, installed.
+
 ## Run tests, linting, packaging checks
 
-```shell
-tox list                # list all Tox environments
-tox run -e ruff         # run code style checks
-tox run -e py           # run tests with your default Python
-tox run -e package      # verify packaging
-tox                     # run all Tox environments
-```
+| Command                   | Description                                                                                                                                                     |
+| ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `nox --list`              | Lists all available Nox sessions, including selected ones                                                                                                       |
+| `nox -s lint`             | Runs code style checks with pre-commit and pre-commit hooks as listed in `.pre-commit-config.yaml`. Accepts posargs to pass additional arguments to the linter. |
+| `nox -s tests`            | Runs tests with your default Python interpreter. Accepts posargs to pass additional arguments and configuration to `pytest`.                                    |
+| `nox -s nightly-tests`    | Similar to `nox -s tests`, except that it runs tests with nightly versions of dependencies (NumPy, SciPy, etc.).                                                |
+| `nox -s validate-package` | Builds a source distribution and a wheel using `pypa/build` and checks the package with `twine` in strict mode.                                                 |
+| `nox`                     | Runs all selected sessions, as listed in `nox.options.sessions` in `noxfile.py`.                                                                                |
+
+Additionally, `nox` supports tags to run specific sessions, e.g., `nox --tags tests` runs all sessions tagged with `tests`.
 
 Make sure all tests pass before you push your changes to GitHub.
 GH Actions will run the tests across all supported Python versions.
 
-## Using arguments (reformat, upload package, help)
-
-You can use additional arguments for the tools called by Tox by
-separating them from the Tox arguments by a double-dash `--`, e.g.
-
-```shell
-tox run -e ruff -- autograd/core.py --show-source
-tox run -e ruff -- autograd/core.py --fix
-```
+## Using positional arguments (reformat, upload package, help)
 
-```shell
-tox run -e package -- upload
-```
+You can use additional arguments for the tools (`pytest`, `pre-commit`, etc.) called by Nox by
+separating them from the Nox arguments by a double-hyphen `--`, e.g.,
 
-```shell
-tox run -e py -- --help
-```
+- `nox -s tests -- --tests/test_tuple.py` runs just the tests listed `tests/test_tuple.py`.
+- `nox -s lint -- --fix` runs the linter with the `--fix` flag.
+- and so on.

noxfile.py

@@ -0,0 +1,60 @@
+import platform
+
+import nox
+
+NIGHTLY_INDEX_URL = "https://pypi.anaconda.org/scientific-python-nightly-wheels/simple"
+UV_NIGHTLY_ENV_VARS = {
+    "UV_INDEX_URL": NIGHTLY_INDEX_URL,
+    "UV_PRERELEASE": "allow",
+    "UV_INDEX_STRATEGY": "first-index",
+    "UV_NO_CACHE": "true"
+}
+
+nox.needs_version = ">=2024.4.15"
+nox.options.default_venv_backend = "uv|virtualenv"
+nox.options.reuse_existing_virtualenvs = False
+nox.options.error_on_external_run = True
+# nox.options.sessions = ["lint", "validate-package", "tests"]
+nox.options.sessions = ["tests"]
+
+
+@nox.session(name="validate-package")
+def check(session):
+    """Build source distribution, wheel, and check their metadata"""
+    session.install("build", "twine", silent=False)
+    session.run("python", "-m", "build")
+    session.run("twine", "check", "--strict", "dist/*")
+
+
+@nox.session(name="tests", tags=["tests"])
+def run_tests(session):
+    """Run unit tests and generate a coverage report"""
+    # SciPy doesn't have wheels on PyPy
+    if platform.python_implementation() == "PyPy":
+        session.install("-e", ".[test]", silent=False)
+    else:
+        session.install("-e", ".[test,scipy]", silent=False)
+    session.run("pytest", "--cov=autograd", "--cov-report=xml", "--cov-append", *session.posargs)
+
+
+# TODO: Replace with pre-commit and pre-commit.ci once
+# https://github.com/HIPS/autograd/pull/634 is merged
+@nox.session(name="lint")
+def ruff(session):
+    """Lightning-fast linting for Python"""
+    session.install("ruff", silent=False)
+    session.run("ruff", "check", ".")
+    session.notify("tests")
+
+
+@nox.session(name="nightly-tests", tags=["tests"])
+def run_nightly_tests(session):
+    """Run tests against nightly versions of dependencies"""
+    session.install("-e", ".[test]", silent=False)
+    # SciPy doesn't have wheels on PyPy
+    if platform.python_implementation() == "PyPy":
+        session.install("numpy", "--upgrade", silent=False, env=UV_NIGHTLY_ENV_VARS)
+    else:
+        session.install(
+            "numpy", "scipy", "--upgrade", silent=False, env=UV_NIGHTLY_ENV_VARS)
+    session.run("pytest", "--cov=autograd", "--cov-report=xml", "--cov-append", *session.posargs)

pyproject.toml

@@ -55,6 +55,11 @@ Source = "https://github.com/HIPS/autograd"
 scipy = [
   "scipy",
 ]
+test = [
+  "pytest",
+  "pytest-cov",
+  "pytest-xdist",
+]
 
 [tool.coverage.run]
 source = ["autograd"]
@@ -63,7 +68,9 @@ source = ["autograd"]
 show_missing = true
 
 [tool.pytest.ini_options]
-addopts = "--color=yes --junitxml=junit-report.xml"
+required_plugins = ["pytest-cov", "pytest-xdist"]
+# TODO: generate HTML report, upload to CodeCov
+addopts = "--color=yes -sra -n auto --cov=autograd --cov-report=xml --cov-report=term"
 
 [tool.ruff]
 extend-exclude = []

tests/conftest.py

@@ -0,0 +1,7 @@
+import numpy as np
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def random_seed():
+    np.random.seed(42)