refactor dependency_resolver.py to use click

python/pip_install/requirements.bzl

@@ -80,12 +80,19 @@ def compile_pip_requirements(
     args = [
         loc.format(requirements_in),
         loc.format(requirements_txt),
-        # String None is a placeholder for argv ordering.
-        loc.format(requirements_linux) if requirements_linux else "None",
-        loc.format(requirements_darwin) if requirements_darwin else "None",
-        loc.format(requirements_windows) if requirements_windows else "None",
         "//%s:%s.update" % (native.package_name(), name),
-    ] + extra_args
+    ]
+
+    if requirements_linux:
+        args.append("--requirements-linux={}".format(loc.format(requirements_linux)))
+
+    if requirements_darwin:
+        args.append("--requirements-darwin={}".format(loc.format(requirements_darwin)))
+
+    if requirements_windows:
+        args.append("--requirements-windows={}".format(loc.format(requirements_windows)))
+
+    args.extend(extra_args)
 
     deps = [
         requirement("build"),

python/pip_install/tools/dependency_resolver/dependency_resolver.py

@@ -19,7 +19,9 @@
 import shutil
 import sys
 from pathlib import Path
+from typing import Optional, Tuple
 
+import click
 import piptools.writer as piptools_writer
 from piptools.scripts.compile import cli
 
@@ -77,24 +79,25 @@ def _locate(bazel_runfiles, file):
     return bazel_runfiles.Rlocation(file)
 
 
-if __name__ == "__main__":
-    if len(sys.argv) < 4:
-        print(
-            "Expected at least two arguments: requirements_in requirements_out",
-            file=sys.stderr,
-        )
-        sys.exit(1)
-
-    parse_str_none = lambda s: None if s == "None" else s
+@click.command(context_settings={"ignore_unknown_options": True})
+@click.argument("requirements_in")
+@click.argument("requirements_txt")
+@click.argument("update_target_label")
+@click.option("--requirements-linux")
+@click.option("--requirements-darwin")
+@click.option("--requirements-windows")
+@click.argument("extra_args", nargs=-1, type=click.UNPROCESSED)
+def main(
+    requirements_in: str,
+    requirements_txt: str,
+    update_target_label: str,
+    requirements_linux: Optional[str],
+    requirements_darwin: Optional[str],
+    requirements_windows: Optional[str],
+    extra_args: Tuple[str, ...],
+) -> None:
     bazel_runfiles = runfiles.Create()
 
-    requirements_in = sys.argv.pop(1)
-    requirements_txt = sys.argv.pop(1)
-    requirements_linux = parse_str_none(sys.argv.pop(1))
-    requirements_darwin = parse_str_none(sys.argv.pop(1))
-    requirements_windows = parse_str_none(sys.argv.pop(1))
-    update_target_label = sys.argv.pop(1)
-
     resolved_requirements_in = _locate(bazel_runfiles, requirements_in)
     resolved_requirements_txt = _locate(bazel_runfiles, requirements_txt)
 
@@ -123,6 +126,8 @@ def _locate(bazel_runfiles, file):
     os.environ["LC_ALL"] = "C.UTF-8"
     os.environ["LANG"] = "C.UTF-8"
 
+    argv = []
+
     UPDATE = True
     # Detect if we are running under `bazel test`.
     if "TEST_TMPDIR" in os.environ:
@@ -131,8 +136,7 @@ def _locate(bazel_runfiles, file):
         # to the real user cache, Bazel sandboxing makes the file read-only
         # and we fail.
         # In theory this makes the test more hermetic as well.
-        sys.argv.append("--cache-dir")
-        sys.argv.append(os.environ["TEST_TMPDIR"])
+        argv.append(f"--cache-dir={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"
@@ -148,15 +152,14 @@ def _locate(bazel_runfiles, file):
     os.environ["CUSTOM_COMPILE_COMMAND"] = update_command
     os.environ["PIP_CONFIG_FILE"] = os.getenv("PIP_CONFIG_FILE") or os.devnull
 
-    sys.argv.append("--generate-hashes")
-    sys.argv.append("--output-file")
-    sys.argv.append(requirements_txt_relative if UPDATE else requirements_out)
-    sys.argv.append(
+    argv.append("--generate-hashes")
+    argv.append(f"--output-file={requirements_txt_relative if UPDATE else requirements_out}")
+    argv.append(
         requirements_in_relative
         if Path(requirements_in_relative).exists()
         else resolved_requirements_in
     )
-    print(sys.argv)
+    argv.extend(extra_args)
 
     if UPDATE:
         print("Updating " + requirements_txt_relative)
@@ -172,7 +175,7 @@ def _locate(bazel_runfiles, file):
                         resolved_requirements_txt, requirements_txt_tree
                     )
                 )
-        cli()
+        cli(argv)
         requirements_txt_relative_path = Path(requirements_txt_relative)
         content = requirements_txt_relative_path.read_text()
         content = content.replace(absolute_path_prefix, "")
@@ -181,7 +184,7 @@ def _locate(bazel_runfiles, file):
         # cli will exit(0) on success
         try:
             print("Checking " + requirements_txt)
-            cli()
+            cli(argv)
             print("cli() should exit", file=sys.stderr)
             sys.exit(1)
         except SystemExit as e:
@@ -221,3 +224,7 @@ def _locate(bazel_runfiles, file):
                     file=sys.stderr,
                 )
                 sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()