fix: use tensorflow_docs.tools.nbfmt for notebook checks (#1274)

The checker called a non-existent tensorflow_docs.tools.nbformat module.
Align with Cirq by using nbfmt, regenerate lockfiles on current main, and
format all tracked notebooks so CI can pass.

Author: rosspeili

check/nbformat

@@ -16,33 +16,28 @@
 set -euo pipefail
 
 declare -r usage="\
-Usage: ${0##*/} [--help] [--fix]
+Usage: ${0##*/} [--help] [--apply]
 
-Checks the format of Jupyter notebook (.ipynb) files in this repository using
-the TensorFlow Docs notebook format checker (tensorflow_docs.tools.nbformat).
+Checks Jupyter notebook (.ipynb) format using tensorflow_docs.tools.nbfmt.
 
-With no arguments, notebooks are checked but not modified. Exits with a
-non-zero status code if any notebook fails the format check.
-
-With --fix, the notebooks are reformatted in place and the script exits with
-status 0 if all notebooks were successfully processed.
+Without --apply, exits 0 if all notebooks are formatted and 1 otherwise.
+With --apply, reformats notebooks in place.
 "
 
 # Get the working directory to the repo root.
 thisdir=$(dirname "${BASH_SOURCE[0]:?}")
 repo_dir=$(git -C "${thisdir}" rev-parse --show-toplevel)
 cd "${repo_dir}"
 
-# Parse arguments.
-opt_fix=0
+opt_apply=0
 for arg in "$@"; do
     case "${arg}" in
         -h | --help)
             echo "${usage}"
             exit 0
             ;;
-        --fix)
-            opt_fix=1
+        --apply)
+            opt_apply=1
             ;;
         *)
             echo "Unknown argument: '${arg}'" >&2
@@ -52,47 +47,32 @@ for arg in "$@"; do
     esac
 done
 
-# Find all tracked .ipynb files in the repository.
-IFS=$'\n' read -r -d '' -a notebooks < <(
-    git ls-files '*.ipynb'
-    printf "\0"
-)
-
-if [[ "${#notebooks[@]}" -eq 0 ]]; then
-    echo "No .ipynb notebook files found."
-    exit 0
-fi
-
-echo "Found ${#notebooks[@]} notebook(s)."
-
-# Build the tensorflow_docs.tools.nbformat argument list.
-declare -a tf_args=()
-if (( opt_fix )); then
-    echo "Running notebook format fixer..."
-    tf_args=( "--fix" )
-else
-    echo "Running notebook format checker..."
+if ! python -c "import tensorflow_docs.tools.nbfmt"; then
+    echo "tensorflow-docs is not installed." >&2
+    exit 2
 fi
 
-declare -a errors=()
-for notebook in "${notebooks[@]}"; do
-    echo "  Checking: ${notebook}"
-    if ! python -m tensorflow_docs.tools.nbformat "${tf_args[@]}" "${notebook}"; then
-        errors+=( "${notebook}" )
-    fi
-done
+declare -a notebook_paths=(
+    docs
+    src/openfermion/resource_estimates/pbc/notebooks
+)
 
-echo
+my_nbfmt() {
+    python -m tensorflow_docs.tools.nbfmt --indent=1 "$@"
+}
 
-if [[ "${#errors[@]}" -eq 0 ]]; then
-    echo "All notebooks passed the format check."
-    exit 0
+exit_code=0
+if output=$(my_nbfmt --test "${notebook_paths[@]}" 2>&1); then
+    echo "Notebooks are formatted."
+elif (( opt_apply )); then
+    my_nbfmt "${notebook_paths[@]}" &&
+    echo "Reformatted notebooks." ||
+    exit_code=$?
 else
-    echo "The following notebooks failed the format check:" >&2
-    printf "  %s\n" "${errors[@]}" >&2
-    if (( ! opt_fix )); then
-        echo "" >&2
-        echo "Run 'check/nbformat --fix' to fix the notebooks automatically." >&2
-    fi
-    exit 1
+    echo "The following notebooks require formatting:" >&2
+    grep '^-' <<<"${output}" >&2 || echo "${output}" >&2
+    echo "Run 'check/nbformat --apply' to fix notebooks." >&2
+    exit_code=1
 fi
+
+exit "${exit_code}"

dev_tools/requirements/envs/dev.env.txt

@@ -1,24 +1,24 @@
 #
-# This file is autogenerated by pip-compile with Python 3.10
+# This file is autogenerated by pip-compile with Python 3.13
 # by the following command:
 #
 #    pip-compile --output-file=envs/dev.env.txt deps/format.txt deps/mypy.txt deps/pip-tools.txt deps/pylint.txt deps/pytest.txt deps/resource_estimates_runtime.txt deps/runtime.txt deps/shellcheck.txt
 #
+absl-py==2.4.0
+    # via tensorflow-docs
 ase==3.28.0
     # via -r deps/resource_estimates_runtime.txt
 ast-serialize==0.5.0
     # via mypy
+astor==0.8.1
+    # via tensorflow-docs
 astroid==3.3.11
     # via pylint
 attrs==26.1.0
     # via
     #   cirq-core
     #   jsonschema
     #   referencing
-backports-asyncio-runner==1.2.0 ; python_full_version < "3.11"
-    # via
-    #   -r deps/runtime.txt
-    #   pytest-asyncio
 black==26.5.1
     # via -r deps/format.txt
 build==1.5.0
@@ -33,6 +33,13 @@ click==8.4.1
     # via
     #   black
     #   pip-tools
+colorama==0.4.6
+    # via
+    #   build
+    #   click
+    #   pylint
+    #   pytest
+    #   tqdm
 contourpy==1.3.2
     # via
     #   -r deps/runtime.txt
@@ -47,8 +54,6 @@ dill==0.4.1
     # via pylint
 duet==0.2.9
     # via cirq-core
-exceptiongroup==1.3.1
-    # via pytest
 execnet==2.1.2
     # via pytest-xdist
 fastjsonschema==2.21.2
@@ -71,6 +76,8 @@ jaxlib==0.4.38
     # via
     #   -r deps/resource_estimates_runtime.txt
     #   jax
+jinja2==3.1.6
+    # via tensorflow-docs
 jsonschema==4.26.0
     # via nbformat
 jsonschema-specifications==2025.9.1
@@ -81,6 +88,8 @@ kiwisolver==1.5.0
     # via matplotlib
 librt==0.11.0
     # via mypy
+markupsafe==3.0.3
+    # via jinja2
 matplotlib==3.10.9
     # via
     #   ase
@@ -100,7 +109,9 @@ mypy-extensions==1.1.0
     #   black
     #   mypy
 nbformat==5.10.4
-    # via -r deps/pytest.txt
+    # via
+    #   -r deps/pytest.txt
+    #   tensorflow-docs
 networkx==3.4.2
     # via
     #   -r deps/runtime.txt
@@ -152,6 +163,8 @@ pluggy==1.6.0
     # via
     #   pytest
     #   pytest-cov
+protobuf==7.35.0
+    # via tensorflow-docs
 pubchempy==1.0.5
     # via -r deps/runtime.txt
 pygments==2.20.0
@@ -192,6 +205,8 @@ pytokens==0.4.1
     # via black
 pytz==2026.2
     # via pandas
+pyyaml==6.0.3
+    # via tensorflow-docs
 referencing==0.37.0
     # via
     #   jsonschema
@@ -220,15 +235,8 @@ sympy==1.14.0
     # via
     #   -r deps/runtime.txt
     #   cirq-core
-tomli==2.4.1
-    # via
-    #   black
-    #   build
-    #   coverage
-    #   mypy
-    #   pip-tools
-    #   pylint
-    #   pytest
+tensorflow-docs==2025.2.19.33219
+    # via -r deps/pytest.txt
 tomlkit==0.15.0
     # via pylint
 tqdm==4.67.3
@@ -247,13 +255,8 @@ types-setuptools==82.0.0.20260518
     # via -r deps/mypy.txt
 typing-extensions==4.15.0
     # via
-    #   astroid
-    #   black
     #   cirq-core
-    #   exceptiongroup
     #   mypy
-    #   pytest-asyncio
-    #   referencing
 tzdata==2026.2
     # via pandas
 urllib3==2.7.0

dev_tools/requirements/envs/pytest.env.txt

@@ -1,20 +1,23 @@
 #
-# This file is autogenerated by pip-compile with Python 3.10
+# This file is autogenerated by pip-compile with Python 3.13
 # by the following command:
 #
 #    pip-compile --constraint=envs/dev.env.txt --output-file=envs/pytest.env.txt deps/pytest.txt deps/runtime.txt
 #
+absl-py==2.4.0
+    # via
+    #   -c envs/dev.env.txt
+    #   tensorflow-docs
+astor==0.8.1
+    # via
+    #   -c envs/dev.env.txt
+    #   tensorflow-docs
 attrs==26.1.0
     # via
     #   -c envs/dev.env.txt
     #   cirq-core
     #   jsonschema
     #   referencing
-backports-asyncio-runner==1.2.0 ; python_full_version < "3.11"
-    # via
-    #   -c envs/dev.env.txt
-    #   -r deps/runtime.txt
-    #   pytest-asyncio
 certifi==2026.5.20
     # via
     #   -c envs/dev.env.txt
@@ -27,6 +30,11 @@ cirq-core==1.5.0
     # via
     #   -c envs/dev.env.txt
     #   -r deps/runtime.txt
+colorama==0.4.6
+    # via
+    #   -c envs/dev.env.txt
+    #   pytest
+    #   tqdm
 contourpy==1.3.2
     # via
     #   -c envs/dev.env.txt
@@ -48,10 +56,6 @@ duet==0.2.9
     # via
     #   -c envs/dev.env.txt
     #   cirq-core
-exceptiongroup==1.3.1
-    # via
-    #   -c envs/dev.env.txt
-    #   pytest
 execnet==2.1.2
     # via
     #   -c envs/dev.env.txt
@@ -76,6 +80,10 @@ iniconfig==2.3.0
     # via
     #   -c envs/dev.env.txt
     #   pytest
+jinja2==3.1.6
+    # via
+    #   -c envs/dev.env.txt
+    #   tensorflow-docs
 jsonschema==4.26.0
     # via
     #   -c envs/dev.env.txt
@@ -92,6 +100,10 @@ kiwisolver==1.5.0
     # via
     #   -c envs/dev.env.txt
     #   matplotlib
+markupsafe==3.0.3
+    # via
+    #   -c envs/dev.env.txt
+    #   jinja2
 matplotlib==3.10.9
     # via
     #   -c envs/dev.env.txt
@@ -104,6 +116,7 @@ nbformat==5.10.4
     # via
     #   -c envs/dev.env.txt
     #   -r deps/pytest.txt
+    #   tensorflow-docs
 networkx==3.4.2
     # via
     #   -c envs/dev.env.txt
@@ -142,6 +155,10 @@ pluggy==1.6.0
     #   -c envs/dev.env.txt
     #   pytest
     #   pytest-cov
+protobuf==7.35.0
+    # via
+    #   -c envs/dev.env.txt
+    #   tensorflow-docs
 pubchempy==1.0.5
     # via
     #   -c envs/dev.env.txt
@@ -192,6 +209,10 @@ pytz==2026.2
     # via
     #   -c envs/dev.env.txt
     #   pandas
+pyyaml==6.0.3
+    # via
+    #   -c envs/dev.env.txt
+    #   tensorflow-docs
 referencing==0.37.0
     # via
     #   -c envs/dev.env.txt
@@ -224,11 +245,10 @@ sympy==1.14.0
     #   -c envs/dev.env.txt
     #   -r deps/runtime.txt
     #   cirq-core
-tomli==2.4.1
+tensorflow-docs==2025.2.19.33219
     # via
     #   -c envs/dev.env.txt
-    #   coverage
-    #   pytest
+    #   -r deps/pytest.txt
 tqdm==4.67.3
     # via
     #   -c envs/dev.env.txt
@@ -242,9 +262,6 @@ typing-extensions==4.15.0
     # via
     #   -c envs/dev.env.txt
     #   cirq-core
-    #   exceptiongroup
-    #   pytest-asyncio
-    #   referencing
 tzdata==2026.2
     # via
     #   -c envs/dev.env.txt