Switch to pyproject.toml; remove setuptools dependency

Also, switch from flake8 to Ruff

Author: NeilGirdhar

.github/workflows/ci.yml

@@ -18,7 +18,7 @@ concurrency:
 
 jobs:
   build-and-test:
-    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }} jax=${{ matrix.jax-version}}"
+    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }} jax=${{ matrix.jax-version }}"
     runs-on: "${{ matrix.os }}"
 
     strategy:
@@ -29,7 +29,7 @@ jobs:
         include:
           - python-version: "3.10"
             os: "ubuntu-latest"
-            jax-version: "0.4.27"  # Keep this in sync with version in requirements.txt
+            jax-version: "0.4.27"  # Keep this in sync with version in pyproject.toml
           - python-version: "3.13"
             os: "ubuntu-latest"
             jax-version: "nightly"
@@ -38,9 +38,22 @@ jobs:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
         with:
-            python-version: "${{ matrix.python-version }}"
-            cache: "pip"
-            cache-dependency-path: '**/requirements*.txt'
+          python-version: "${{ matrix.python-version }}"
+
+      - name: Compute pyproject.toml hash
+        id: deps
+        run: |
+          echo "hash=$(shasum -a 256 pyproject.toml | cut -d' ' -f1)" >> $GITHUB_OUTPUT
+
+      - name: Cache pip based on pyproject.toml
+        uses: actions/cache@v3
+        with:
+          path: ~/.cache/pip
+          key: pip-${{ runner.os }}-${{ matrix.python-version }}-${{ steps.deps.outputs.hash }}
+
+      - name: Install build dependencies
+        run: pip install .[dev,test]
+
       - name: Run CI tests
         run: JAX_VERSION="${{ matrix.jax-version }}" bash test.sh
         shell: bash

MANIFEST.in

@@ -0,0 +1,84 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "chex"
+dynamic = ["version"]
+description = "Chex: Testing made fun, in JAX!"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.9"
+authors = [
+    { name = "DeepMind", email = "[email protected]" },
+]
+keywords = [
+    "debugging",
+    "jax",
+    "learning",
+    "machine",
+    "python",
+    "testing",
+]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Environment :: Console",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Software Development :: Testing :: Mocking",
+    "Topic :: Software Development :: Testing :: Unit",
+]
+dependencies = [
+    "absl-py>=0.9.0",
+    "jax>=0.4.27",
+    "jaxlib>=0.4.27",
+    "numpy>=1.24.1",
+    "toolz>=0.9.0",
+    "typing_extensions>=4.2.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "ruff",
+    "pytest-xdist",
+    "pylint",
+    "pylint-exit",
+]
+test = [
+    "cloudpickle==3.1.0",
+    "dm-tree>=0.1.9",
+]
+docs = [
+    "sphinx>=6.0.0",
+    "sphinx-book-theme>=1.0.1",
+    "sphinxcontrib-katex",
+]
+
+[project.urls]
+Homepage = "https://github.com/deepmind/chex"
+
+[tool.hatch.version]
+path = "chex/__init__.py"
+
+[tool.hatch.build.targets.sdist]
+include = [
+    "/chex",
+]
+
+[tool.ruff]
+preview = true
+
+[tool.ruff.lint]
+select = [
+    "E9",
+    "F63",
+    "F7",
+    "F82",
+    "E225",
+    "E251",
+]

pyproject.toml

@@ -27,10 +27,8 @@ source "${VENV_DIR}/bin/activate"
 python --version
 
 # Install dependencies.
-pip install --upgrade pip setuptools wheel
-pip install flake8 pytest-xdist pylint pylint-exit
-pip install -r requirements/requirements.txt
-pip install -r requirements/requirements-test.txt
+pip install --upgrade pip
+pip install .[dev]
 
 # Install the requested JAX version
 if [ "$JAX_VERSION" = "" ]; then
@@ -43,8 +41,8 @@ else
   pip install "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}"
 fi
 
-# Lint with flake8.
-flake8 `find chex -name '*.py' | xargs` --count --select=E9,F63,F7,F82,E225,E251 --show-source --statistics
+# Lint with Ruff.
+ruff check
 
 # Lint with pylint.
 PYLINT_ARGS="-efail -wfail -cfail -rfail"
@@ -62,9 +60,10 @@ pylint --rcfile=.pylintrc `find chex -name '*_test.py' | xargs` -d W0212,E1130,E
 rm .pylintrc
 
 # Build the package.
-python setup.py sdist
+pip install build
+python -m build
 pip wheel --verbose --no-deps --no-clean dist/chex*.tar.gz
-pip install chex*.whl
+pip install chex*.whl --force-reinstall
 
 # Check types with pytype.
 # Note: pytype does not support 3.11 as of 25.06.23
@@ -77,7 +76,7 @@ fi;
 
 # Run tests using pytest.
 # Change directory to avoid importing the package from repo root.
-pip install -r requirements/requirements-test.txt
+pip install .[test]
 cd _testing
 
 # Main tests.
@@ -88,8 +87,7 @@ pytest -n "$(grep -c ^processor /proc/cpuinfo)" --pyargs chex -k "fake_set_n_cpu
 cd ..
 
 # Build Sphinx docs.
-
-pip install -r requirements/requirements-docs.txt
+pip install .[docs]
 cd docs
 make coverage_check
 make html