Improve replace_version tool and development wheels (#6626)

- Make `dev_tools/modules.py replace_version` update also the
  version expected in unit tests

- Skip type check in the new test-only source files

- Update `dev_tools/packaging/produce-package.sh` to use the `replace_version`
  utility when building development wheels

dev_tools/modules.py

@@ -160,16 +160,24 @@ def replace_version(search_dir: Path = _DEFAULT_SEARCH_DIR, *, old: str, new: st
 
     for m in list_modules(search_dir=search_dir, include_parent=True):
         version_file = _find_version_file(search_dir / m.root)
-        content = version_file.read_text("UTF-8")
-        new_content = content.replace(old, new)
-        version_file.write_text(new_content)
+        _rewrite_version(version_file, old, new)
+        version_test = version_file.parent / "_version_test.py"
+        _rewrite_version(version_test, old, new)
 
 
 def _validate_version(new_version: str):
     if not re.match(r"\d+\.\d+\.\d+(\.dev)?", new_version):
         raise ValueError(f"{new_version} is not a valid version number.")
 
 
+def _rewrite_version(version_file: Path, old: str, new: str) -> None:
+    pattern = f"(^[^#]*__version__ ==? )(['\"])({re.escape(old)})(\\2)"
+    repl = f"\\1\\g<2>{new}\\4"
+    content = version_file.read_text("UTF-8")
+    new_content = re.sub(pattern, repl, content, flags=re.MULTILINE)
+    version_file.write_text(new_content)
+
+
 def _find_version_file(top: Path) -> Path:
     for root, _, files in os.walk(str(top)):
         if "_version.py" in files:

dev_tools/modules_test.py

@@ -149,6 +149,7 @@ def test_get_version_on_no_modules():
 def test_get_version_on_inconsistent_version_modules():
     modules.replace_version(search_dir=Path("./mod2"), old="1.2.3.dev", new="1.2.4.dev")
     assert modules.get_version(search_dir=Path("./mod2")) == "1.2.4.dev"
+    assert "1.2.4.dev" in Path("./mod2/pack2/_version_test.py").read_text("UTF-8")
     with pytest.raises(ValueError, match="Versions should be the same, instead:"):
         modules.get_version(search_dir=Path("."))
 
@@ -158,6 +159,8 @@ def test_replace_version(tmpdir_factory):
     assert modules.get_version() == "1.2.3.dev"
     modules.replace_version(old="1.2.3.dev", new="1.2.4.dev")
     assert modules.get_version() == "1.2.4.dev"
+    assert "1.2.4.dev" in Path("./mod1/pack1/_version_test.py").read_text("UTF-8")
+    assert "1.2.4.dev" in Path("./mod2/pack2/_version_test.py").read_text("UTF-8")
 
 
 @chdir(target_dir="dev_tools/modules_test_data")

dev_tools/modules_test_data/mod1/pack1/_version_test.py

@@ -0,0 +1,6 @@
+# pylint: disable=wrong-or-nonexistent-copyright-notice
+import pack1._version  # type: ignore
+
+
+def test_version():
+    assert pack1._version.__version__ == "1.2.3.dev"

dev_tools/modules_test_data/mod2/pack2/_version_test.py

@@ -0,0 +1,6 @@
+# pylint: disable=wrong-or-nonexistent-copyright-notice
+import pack2._version  # type: ignore
+
+
+def test_version():
+    assert pack2._version.__version__ == "1.2.3.dev"

dev_tools/packaging/produce-package.sh

@@ -36,6 +36,13 @@ out_dir=$(realpath "${1}")
 
 SPECIFIED_VERSION="${2}"
 
+# Helper to run dev_tools/modules.py without CIRQ_PRE_RELEASE_VERSION
+# to avoid environment version override in setup.py.
+my_dev_tools_modules() {
+    env -u CIRQ_PRE_RELEASE_VERSION PYTHONPATH=. \
+        python3 dev_tools/modules.py "$@"
+}
+
 # Get the working directory to the repo root.
 cd "$( dirname "${BASH_SOURCE[0]}" )"
 repo_dir=$(git rev-parse --show-toplevel)
@@ -47,21 +54,18 @@ if [ -n "$(git status --short)" ]; then
 fi
 tmp_git_dir=$(mktemp -d "/tmp/produce-package-git.XXXXXXXXXXXXXXXX")
 trap '{ rm -rf "${tmp_git_dir}"; }' EXIT
+echo "Creating pristine repository clone at ${tmp_git_dir}"
+git clone --shared --quiet "${repo_dir}" "${tmp_git_dir}"
 cd "${tmp_git_dir}"
-git init --quiet
-git fetch "${repo_dir}" HEAD --quiet --depth=1
-git checkout FETCH_HEAD -b work --quiet
 if [ -n "${SPECIFIED_VERSION}" ]; then
-    CIRQ_PACKAGES=$(env PYTHONPATH=. python dev_tools/modules.py list --mode package-path)
-    for PROJECT_NAME in $CIRQ_PACKAGES; do
-      echo '__version__ = "'"${SPECIFIED_VERSION}"'"' > "${tmp_git_dir}/${PROJECT_NAME}/_version.py"
-    done
+    CURRENT_VERSION=$(my_dev_tools_modules print_version)
+    my_dev_tools_modules replace_version --old="${CURRENT_VERSION}" --new="${SPECIFIED_VERSION}"
 fi
 
 # Python 3 wheel.
 echo "Producing python 3 package files."
 
-CIRQ_MODULES=$(env PYTHONPATH=. python dev_tools/modules.py list --mode folder --include-parent)
+CIRQ_MODULES=$(my_dev_tools_modules list --mode folder --include-parent)
 
 for m in $CIRQ_MODULES; do
   echo "processing $m/setup.py..."