refactor: Unify invoking pip-compile

Author: brandonchinn178

python/private/pypi/dependency_resolver/dependency_resolver.py

@@ -15,11 +15,12 @@
 "Set defaults for the pip-compile command to run it under Bazel"
 
 import atexit
+import functools
 import os
 import shutil
 import sys
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import click
 import piptools.writer as piptools_writer
@@ -168,6 +169,12 @@ def main(
     )
     argv.extend(extra_args)
 
+    _run_pip_compile = functools.partial(
+        run_pip_compile,
+        argv,
+        srcs_relative=srcs_relative,
+    )
+
     if UPDATE:
         print("Updating " + requirements_file_relative)
 
@@ -187,49 +194,52 @@ def main(
                 atexit.register(
                     lambda: shutil.copy(absolute_output_file, requirements_file_tree)
                 )
-        cli(argv, standalone_mode=False)
+        _run_pip_compile()
         requirements_file_relative_path = Path(requirements_file_relative)
         content = requirements_file_relative_path.read_text()
         content = content.replace(absolute_path_prefix, "")
         requirements_file_relative_path.write_text(content)
     else:
-        # cli will exit(0) on success
-        try:
-            print("Checking " + requirements_file)
-            cli(argv)
-            print("cli() should exit", file=sys.stderr)
+        print("Checking " + requirements_file)
+        _run_pip_compile()
+        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:
+            import difflib
+
+            print("".join(difflib.unified_diff(golden, out)), file=sys.stderr)
+            print(
+                "Lock file out of date. Run '" + update_command + "' to update.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+
+def run_pip_compile(
+    args: List[str],
+    *,
+    srcs_relative: List[str],
+) -> None:
+    try:
+        cli(args, standalone_mode=False)
+    except SystemExit as e:
+        if e.code == 0:
+            return  # shouldn't happen, but just in case
+        elif e.code == 2:
+            print(
+                "pip-compile exited with code 2. This means that pip-compile found "
+                "incompatible requirements or could not find a version that matches "
+                f"the install requirement in one of {srcs_relative}.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        else:
+            print(
+                f"pip-compile unexpectedly exited with code {e.code}.",
+                file=sys.stderr,
+            )
             sys.exit(1)
-        except SystemExit as e:
-            if e.code == 2:
-                print(
-                    "pip-compile exited with code 2. This means that pip-compile found "
-                    "incompatible requirements or could not find a version that matches "
-                    f"the install requirement in one of {srcs_relative}.",
-                    file=sys.stderr,
-                )
-                sys.exit(1)
-            elif e.code == 0:
-                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:
-                    import difflib
-
-                    print("".join(difflib.unified_diff(golden, out)), file=sys.stderr)
-                    print(
-                        "Lock file out of date. Run '"
-                        + update_command
-                        + "' to update.",
-                        file=sys.stderr,
-                    )
-                    sys.exit(1)
-                sys.exit(0)
-            else:
-                print(
-                    f"pip-compile unexpectedly exited with code {e.code}.",
-                    file=sys.stderr,
-                )
-                sys.exit(1)
 
 
 if __name__ == "__main__":