fix: update correct requirements lock file when using os specific lock files (#1123)

Currently the dependency_resolver.py ignores that you give requirement
lock files for different os's, except when checking if the golden file
needs updating. This causes dependecy_resolver.py to update the wrong
lock i.e the non platform specific one if ran in "update mode".

6 files changed, 109 insertions, 26 deletions

diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index b468970..ac24113 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -435,6 +435,12 @@ tasks:
     - "! git diff --exit-code"
     - "bazel run //:requirements.update"
     - "git diff --exit-code"
+    # Make a change to the locked requirements and then assert that //:os_specific_requirements.update does the
+    # right thing.
+    - "echo '' > requirements_lock_linux.txt"
+    - "! git diff --exit-code"
+    - "bazel run //:os_specific_requirements.update"
+    - "git diff --exit-code"
   integration_test_compile_pip_requirements_debian:
     <<: *reusable_build_test_all
     name: compile_pip_requirements integration tests on Debian
@@ -447,6 +453,12 @@ tasks:
     - "! git diff --exit-code"
     - "bazel run //:requirements.update"
     - "git diff --exit-code"
+    # Make a change to the locked requirements and then assert that //:os_specific_requirements.update does the
+    # right thing.
+    - "echo '' > requirements_lock_linux.txt"
+    - "! git diff --exit-code"
+    - "bazel run //:os_specific_requirements.update"
+    - "git diff --exit-code"
   integration_test_compile_pip_requirements_macos:
     <<: *reusable_build_test_all
     name: compile_pip_requirements integration tests on macOS
@@ -459,6 +471,12 @@ tasks:
     - "! git diff --exit-code"
     - "bazel run //:requirements.update"
     - "git diff --exit-code"
+    # Make a change to the locked requirements and then assert that //:os_specific_requirements.update does the
+    # right thing.
+    - "echo '' > requirements_lock_darwin.txt"
+    - "! git diff --exit-code"
+    - "bazel run //:os_specific_requirements.update"
+    - "git diff --exit-code"
   integration_test_compile_pip_requirements_windows:
     <<: *reusable_build_test_all
     name: compile_pip_requirements integration tests on Windows
@@ -471,6 +489,12 @@ tasks:
     - "! git diff --exit-code"
     - "bazel run //:requirements.update"
     - "git diff --exit-code"
+    # Make a change to the locked requirements and then assert that //:os_specific_requirements.update does the
+    # right thing.
+    - "echo '' > requirements_lock_windows.txt"
+    - "! git diff --exit-code"
+    - "bazel run //:os_specific_requirements.update"
+    - "git diff --exit-code"
 
   integration_test_pip_repository_entry_points_ubuntu_min:
     <<: *minimum_supported_version
diff --git a/python/pip_install/tools/dependency_resolver/dependency_resolver.py b/python/pip_install/tools/dependency_resolver/dependency_resolver.py
index 89e3558..ceb20db 100644
--- a/python/pip_install/tools/dependency_resolver/dependency_resolver.py
+++ b/python/pip_install/tools/dependency_resolver/dependency_resolver.py
@@ -95,25 +95,30 @@ if __name__ == "__main__":
     requirements_windows = parse_str_none(sys.argv.pop(1))
     update_target_label = sys.argv.pop(1)
 
+    requirements_file = _select_golden_requirements_file(
+        requirements_txt=requirements_txt, requirements_linux=requirements_linux,
+        requirements_darwin=requirements_darwin, requirements_windows=requirements_windows
+    )
+
     resolved_requirements_in = _locate(bazel_runfiles, requirements_in)
-    resolved_requirements_txt = _locate(bazel_runfiles, requirements_txt)
+    resolved_requirements_file = _locate(bazel_runfiles, requirements_file)
 
     # Files in the runfiles directory has the following naming schema:
     # Main repo: __main__/<path_to_file>
     # External repo: <workspace name>/<path_to_file>
     # We want to strip both __main__ and <workspace name> from the absolute prefix
     # to keep the requirements lock file agnostic.
-    repository_prefix = requirements_txt[: requirements_txt.index("/") + 1]
-    absolute_path_prefix = resolved_requirements_txt[
-        : -(len(requirements_txt) - len(repository_prefix))
+    repository_prefix = requirements_file[: requirements_file.index("/") + 1]
+    absolute_path_prefix = resolved_requirements_file[
+        : -(len(requirements_file) - len(repository_prefix))
     ]
 
     # As requirements_in might contain references to generated files we want to
     # use the runfiles file first. Thus, we need to compute the relative path
     # from the execution root.
     # Note: Windows cannot reference generated files without runfiles support enabled.
-    requirements_in_relative = requirements_in[len(repository_prefix) :]
-    requirements_txt_relative = requirements_txt[len(repository_prefix) :]
+    requirements_in_relative = requirements_in[len(repository_prefix):]
+    requirements_file_relative = requirements_file[len(repository_prefix):]
 
     # Before loading click, set the locale for its parser.
     # If it leaks through to the system setting, it may fail:
@@ -135,11 +140,11 @@ if __name__ == "__main__":
         sys.argv.append(os.environ["TEST_TMPDIR"])
         # Make a copy for pip-compile to read and mutate.
         requirements_out = os.path.join(
-            os.environ["TEST_TMPDIR"], os.path.basename(requirements_txt) + ".out"
+            os.environ["TEST_TMPDIR"], os.path.basename(requirements_file) + ".out"
         )
         # Those two files won't necessarily be on the same filesystem, so we can't use os.replace
         # or shutil.copyfile, as they will fail with OSError: [Errno 18] Invalid cross-device link.
-        shutil.copy(resolved_requirements_txt, requirements_out)
+        shutil.copy(resolved_requirements_file, requirements_out)
 
     update_command = os.getenv("CUSTOM_COMPILE_COMMAND") or "bazel run %s" % (
         update_target_label,
@@ -150,7 +155,7 @@ if __name__ == "__main__":
 
     sys.argv.append("--generate-hashes")
     sys.argv.append("--output-file")
-    sys.argv.append(requirements_txt_relative if UPDATE else requirements_out)
+    sys.argv.append(requirements_file_relative if UPDATE else requirements_out)
     sys.argv.append(
         requirements_in_relative
         if Path(requirements_in_relative).exists()
@@ -159,28 +164,28 @@ if __name__ == "__main__":
     print(sys.argv)
 
     if UPDATE:
-        print("Updating " + requirements_txt_relative)
+        print("Updating " + requirements_file_relative)
         if "BUILD_WORKSPACE_DIRECTORY" in os.environ:
             workspace = os.environ["BUILD_WORKSPACE_DIRECTORY"]
-            requirements_txt_tree = os.path.join(workspace, requirements_txt_relative)
-            # In most cases, requirements_txt will be a symlink to the real file in the source tree.
-            # If symlinks are not enabled (e.g. on Windows), then requirements_txt will be a copy,
+            requirements_file_tree = os.path.join(workspace, requirements_file_relative)
+            # In most cases, requirements_file will be a symlink to the real file in the source tree.
+            # If symlinks are not enabled (e.g. on Windows), then requirements_file will be a copy,
             # and we should copy the updated requirements back to the source tree.
-            if not os.path.samefile(resolved_requirements_txt, requirements_txt_tree):
+            if not os.path.samefile(resolved_requirements_file, requirements_file_tree):
                 atexit.register(
                     lambda: shutil.copy(
-                        resolved_requirements_txt, requirements_txt_tree
+                        resolved_requirements_file, requirements_file_tree
                     )
                 )
         cli()
-        requirements_txt_relative_path = Path(requirements_txt_relative)
-        content = requirements_txt_relative_path.read_text()
+        requirements_file_relative_path = Path(requirements_file_relative)
+        content = requirements_file_relative_path.read_text()
         content = content.replace(absolute_path_prefix, "")
-        requirements_txt_relative_path.write_text(content)
+        requirements_file_relative_path.write_text(content)
     else:
         # cli will exit(0) on success
         try:
-            print("Checking " + requirements_txt)
+            print("Checking " + requirements_file)
             cli()
             print("cli() should exit", file=sys.stderr)
             sys.exit(1)
@@ -194,13 +199,7 @@ if __name__ == "__main__":
                 )
                 sys.exit(1)
             elif e.code == 0:
-                golden_filename = _select_golden_requirements_file(
-                    requirements_txt,
-                    requirements_linux,
-                    requirements_darwin,
-                    requirements_windows,
-                )
-                golden = open(_locate(bazel_runfiles, golden_filename)).readlines()
+                golden = open(_locate(bazel_runfiles, requirements_file)).readlines()
                 out = open(requirements_out).readlines()
                 out = [line.replace(absolute_path_prefix, "") for line in out]
                 if golden != out:
diff --git a/tests/compile_pip_requirements/BUILD.bazel b/tests/compile_pip_requirements/BUILD.bazel
index d6ac008..87ffe70 100644
--- a/tests/compile_pip_requirements/BUILD.bazel
+++ b/tests/compile_pip_requirements/BUILD.bazel
@@ -32,3 +32,33 @@ compile_pip_requirements(
     requirements_in = "requirements.txt",
     requirements_txt = "requirements_lock.txt",
 )
+
+genrule(
+    name = "generate_os_specific_requirements_in",
+    srcs = [],
+    outs = ["requirements_os_specific.in"],
+    cmd = """
+cat > $@ <<EOF
+pip==22.3.0 ; sys_platform == "linux"
+pip==22.2.2 ; sys_platform == "darwin"
+pip==22.2.1 ; sys_platform == "win32"
+EOF
+""",
+)
+
+compile_pip_requirements(
+    name = "os_specific_requirements",
+    data = [
+        "requirements_extra.in",
+        "requirements_os_specific.in",
+    ],
+    extra_args = [
+        "--allow-unsafe",
+        "--resolver=backtracking",
+    ],
+    requirements_darwin = "requirements_lock_darwin.txt",
+    requirements_in = "requirements_os_specific.in",
+    requirements_linux = "requirements_lock_linux.txt",
+    requirements_txt = "requirements_lock.txt",
+    requirements_windows = "requirements_lock_windows.txt",
+)
diff --git a/tests/compile_pip_requirements/requirements_lock_darwin.txt b/tests/compile_pip_requirements/requirements_lock_darwin.txt
new file mode 100644
index 0000000..7b580a2
--- /dev/null
+++ b/tests/compile_pip_requirements/requirements_lock_darwin.txt
@@ -0,0 +1,10 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    bazel run //:os_specific_requirements.update
+#
+pip==22.2.2 ; sys_platform == "darwin" \
+    --hash=sha256:3fd1929db052f056d7a998439176d3333fa1b3f6c1ad881de1885c0717608a4b \
+    --hash=sha256:b61a374b5bc40a6e982426aede40c9b5a08ff20e640f5b56977f4f91fed1e39a
+    # via -r requirements_os_specific.in
diff --git a/tests/compile_pip_requirements/requirements_lock_linux.txt b/tests/compile_pip_requirements/requirements_lock_linux.txt
new file mode 100644
index 0000000..54eca17
--- /dev/null
+++ b/tests/compile_pip_requirements/requirements_lock_linux.txt
@@ -0,0 +1,10 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    bazel run //:os_specific_requirements.update
+#
+pip==22.3 ; sys_platform == "linux" \
+    --hash=sha256:1daab4b8d3b97d1d763caeb01a4640a2250a0ea899e257b1e44b9eded91e15ab \
+    --hash=sha256:8182aec21dad6c0a49a2a3d121a87cd524b950e0b6092b181625f07ebdde7530
+    # via -r requirements_os_specific.in
diff --git a/tests/compile_pip_requirements/requirements_lock_windows.txt b/tests/compile_pip_requirements/requirements_lock_windows.txt
new file mode 100644
index 0000000..5803d8e
--- /dev/null
+++ b/tests/compile_pip_requirements/requirements_lock_windows.txt
@@ -0,0 +1,10 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    bazel run //:os_specific_requirements.update
+#
+pip==22.2.1 ; sys_platform == "win32" \
+    --hash=sha256:0bbbc87dfbe6eed217beff0021f8b7dea04c8f4a0baa9d31dc4cff281ffc5b2b \
+    --hash=sha256:50516e47a2b79e77446f0d05649f0d53772c192571486236b1905492bfc24bac
+    # via -r requirements_os_specific.in