chore(build): migrate to pyproject.toml for packaging

.github/workflows/ci.yml

@@ -18,9 +18,11 @@ jobs:
 
     steps:
       - uses: "actions/checkout@v2"
-      - uses: "actions/setup-python@v1"
+      - uses: "actions/setup-python@v4"
         with:
             python-version: "${{ matrix.python-version }}"
+            cache: "pip"
+            cache-dependency-path: "pyproject.toml"
       - name: Run CI tests
         run: bash test.sh
         shell: bash

.github/workflows/pypi-publish.yml

@@ -16,11 +16,12 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel twine build
     - name: Check consistency between the package version and release tag
       run: |
+        pip install .
         RELEASE_VER=${GITHUB_REF#refs/*/}
-        PACKAGE_VER="v`python setup.py --version`"
+        PACKAGE_VER="v`python -c 'import dm_pix; print(dm_pix.__version__)'`"
         if [ $RELEASE_VER != $PACKAGE_VER ]
         then
           echo "package ver. ($PACKAGE_VER) != release ver. ($RELEASE_VER)"; exit 1
@@ -30,5 +31,5 @@ jobs:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
         TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
       run: |
-        python setup.py sdist bdist_wheel
+        python -m build
         twine upload dist/*

README.md

@@ -23,7 +23,7 @@ through [`jax.jit`][jit], [`jax.vmap`][vmap] and [`jax.pmap`][pmap].
 PIX is written in pure Python, but depends on C++ code via JAX.
 
 Because JAX installation is different depending on your CUDA version, PIX does
-not list JAX as a dependency in [`requirements.txt`], although it is technically
+not list JAX as a dependency in [`pyproject.toml`], although it is technically
 listed for reference, but commented.
 
 First, follow [JAX installation instructions] to install JAX with the relevant
@@ -110,7 +110,7 @@ If you already have PIX installed, you just need to install some extra
 dependencies and run `pytest` as follows:
 
 ```bash
-$ pip install -r requirements_tests.txt
+$ pip install -e ".[test]"
 $ python -m pytest [-n <NUMCPUS>] dm_pix
 ```
 

pyproject.toml

@@ -0,0 +1,67 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "dm_pix"
+dynamic = ["version"]
+description = 'PIX is an image processing library in JAX, for JAX.'
+readme = "README.md"
+license = { file = "LICENSE" }
+requires-python = ">=3.8"
+authors = [
+    {name = "Google DeepMind", email = "[email protected]"},
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Image Processing",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+    "chex>=0.0.6",
+    # jax>=0.2.17
+    # jaxlib>=0.1.69
+]
+
+[project.optional-dependencies]
+extras = [
+    "jax>=0.2.17",
+    "jaxlib>=0.1.69",
+]
+
+test = [
+    "scipy",
+    "tensorflow",
+    "pytest-xdist",
+]
+
+docs = [
+    "sphinx==4.5.0",
+    "sphinx_rtd_theme==1.0.0",
+    "sphinxcontrib-katex==0.9.0",
+    "sphinxcontrib-bibtex==2.4.2",
+    "sphinx-autodoc-typehints==1.11.1",
+    "IPython==8.10.0",
+    "ipykernel==5.3.4",
+    "pandoc==1.1.0",
+    "myst_nb==0.13.1",
+    "docutils==0.16",
+    "matplotlib==3.5.0",
+]
+
+examples = [
+    "Pillow==10.3.0",
+]
+
+[tool.setuptools.packages.find]
+include=["dm_pix/py.typed"]
+exclude = ["*_test.py", "examples"]

readthedocs.yml

@@ -15,7 +15,9 @@ sphinx:
 
 python:
   install:
-    - requirements: requirements_docs.txt
-    - requirements: requirements.txt
-    - method: setuptools
+    - method: pip
       path: .
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs

test.sh

@@ -24,16 +24,15 @@ python -m venv "${VENV_DIR}"
 source "${VENV_DIR}/bin/activate"
 python --version
 
-# Install JAX.
+# Update pip + setuptools and install dependencies specified in pyproject.toml
 python -m pip install --upgrade pip setuptools
-python -m pip install -r requirements.txt
-python -c 'import jax; print(jax.__version__)'
-
-# Run setup.py, this installs the python dependencies
 python -m pip install .
 
+# print jax version
+python -c 'import jax; print(jax.__version__)'
+
 # Python test dependencies.
-python -m pip install -r requirements_tests.txt
+python -m pip install -e ".[test]"
 
 # CPU count on macos or linux
 if [ "$(uname)" == "Darwin" ]; then