Tidy up test.sh and make it clean up properly.

.gitignore

@@ -11,6 +11,7 @@ _testing/
 docs/_autosummary
 docs/_collections
 docs/modules/generated
+docs/sg_execution_times.rst
 
 # Mac OS
 .DS_Store

README.md

@@ -95,7 +95,7 @@ git clone https://github.com/google-deepmind/optax.git
 To run the tests, please execute the following script.
 
 ```sh
-sh ./test.sh
+sh test.sh
 ```
 
 ### Documentation
@@ -106,7 +106,7 @@ pip install -e ".[docs]"
 ```
 Then, execute the following.
 ```sh
-cd docs/
+cd docs
 make html
 ```
 

test.sh

@@ -13,76 +13,79 @@
 # limitations under the License.
 # ==============================================================================
 
-# Runs CI tests on a local machine.
-set -euo pipefail
+function cleanup {
+  deactivate
+  rm -r "${TEMP_DIR}"
+}
+trap cleanup EXIT
+
+REPO_DIR=$(pwd)
+TEMP_DIR=$(mktemp --directory)
+
+set -o errexit
+set -o nounset
+set -o pipefail
 
 # Install deps in a virtual env.
-rm -rf _testing
-rm -rf dist/
-rm -rf *.whl
-rm -rf .pytype
-mkdir -p _testing
-readonly VENV_DIR="$(mktemp -d $(pwd)/_testing/optax-env.XXXXXXXX)"
-# in the unlikely case in which there was something in that directory
-python3 -m venv "${VENV_DIR}"
-source "${VENV_DIR}/bin/activate"
+python3 -m venv "${TEMP_DIR}/venv"
+source "${TEMP_DIR}/venv/bin/activate"
 
 # Install dependencies.
-python3 -m pip install -q --upgrade pip setuptools wheel
-python3 -m pip install -q flake8 pytest-xdist pylint pylint-exit
-python3 -m pip install -q -e ".[test, examples]"
+python3 -m pip install --quiet --upgrade pip setuptools wheel
+python3 -m pip install --quiet --upgrade flake8 pytest-xdist pylint pylint-exit
+python3 -m pip install --quiet --editable ".[test, examples]"
 
 # Dp-accounting specifies exact minor versions as requirements which sometimes
 # become incompatible with other libraries optax needs. We therefore install
 # dependencies for dp-accounting manually.
 # TODO(b/239416992): Remove this workaround if dp-accounting switches to minimum
 # version requirements.
-python3 -m pip install -q -e ".[dp-accounting]"
-python3 -m pip install -q "dp-accounting>=0.1.1" --no-deps
+python3 -m pip install --quiet --editable ".[dp-accounting]"
+python3 -m pip install --quiet --no-deps "dp-accounting>=0.1.1"
 
 # Install the requested JAX version
 if [ -z "${JAX_VERSION-}" ]; then
   : # use version installed in requirements above
 elif [ "$JAX_VERSION" = "newest" ]; then
-  python3 -m pip install -U jax jaxlib
+  python3 -m pip install --quiet --upgrade jax jaxlib
 else
-  python3 -m pip install "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}"
+  python3 -m pip install --quiet "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}"
 fi
 
 # Ensure optax was not installed by one of the dependencies above,
 # since if it is, the tests below will be run against that version instead of
 # the branch build.
-python3 -m pip uninstall -q -y optax || true
+python3 -m pip uninstall --quiet --yes optax
 
 # Lint with flake8.
-flake8 $(find optax examples -name '*.py' | xargs) --select=E9,F63,F7,F82,E225,E251 --show-source --statistics
+python3 -m flake8 --select=E9,F63,F7,F82,E225,E251 --show-source --statistics
 
 # Lint with pylint.
 PYLINT_ARGS="-efail -wfail -cfail -rfail"
 # Append specific config lines.
 # Lint modules and tests separately.
-pylint --rcfile=.pylintrc $(find optax examples -name '*.py' | grep -v 'test.py' | xargs) -d E1102 || pylint-exit $PYLINT_ARGS $?
+python3 -m pylint --rcfile=.pylintrc $(find optax -name '*.py' | grep -v 'test.py' | xargs) -d E1102 || pylint-exit $PYLINT_ARGS $?
 # Disable protected-access warnings for tests.
-pylint --rcfile=.pylintrc $(find optax examples -name '*_test.py' | xargs) -d W0212,E1102 || pylint-exit $PYLINT_ARGS $?
+python3 -m pylint --rcfile=.pylintrc $(find optax -name '*_test.py' | xargs) -d W0212,E1102 || pylint-exit $PYLINT_ARGS $?
 
 # Build the package.
-python3 -m pip install build
+python3 -m pip install --quiet build
 python3 -m build
-python3 -m pip wheel --verbose --no-deps --no-clean dist/optax*.tar.gz
-python3 -m pip install optax*.whl
+python3 -m pip wheel --no-deps dist/optax-*.tar.gz --wheel-dir "${TEMP_DIR}"
+python3 -m pip install --quiet "${TEMP_DIR}/optax-"*.whl
 
 # Check types with pytype.
-python3 -m pip install -q pytype
-pytype $(find optax/_src examples optax/contrib -name '*.py' | xargs) -k -d import-error
+python3 -m pip install --quiet pytype
+pytype "optax" --keep-going --disable import-error
 
 # Run tests using pytest.
 # Change directory to avoid importing the package from repo root.
-cd _testing
-python3 -m pytest -n auto --pyargs optax
-cd ..
+cd "${TEMP_DIR}"
+python3 -m pytest --numprocesses auto --pyargs optax
+cd "${REPO_DIR}"
 
 # Build Sphinx docs.
-python3 -m pip install -q -e ".[docs]"
+python3 -m pip install --quiet --editable ".[docs]"
 # NOTE(vroulet) We have dependencies issues:
 # tensorflow > 2.13.1 requires ml-dtypes <= 0.3.2
 # but jax requires ml-dtypes >= 0.4.0
@@ -92,15 +95,11 @@ python3 -m pip install -q -e ".[docs]"
 # bug (which issues conflict warnings but runs fine).
 # A long term solution is probably to fully remove tensorflow from our
 # dependencies.
-python3 -m pip install --upgrade -v typing_extensions
-cd docs && make html
+python3 -m pip install --upgrade --verbose typing_extensions
+cd docs
+make html
 # run doctests
 make doctest
 cd ..
 
-# cleanup
-rm -rf _testing
-
-set +u
-deactivate
 echo "All tests passed. Congrats!"