From e7d98918cc9c0d45a60f9be80177a4849a781baa Mon Sep 17 00:00:00 2001
From: Alexey Tereshenkov <50622389+AlexTereshenkov@users.noreply.github.com>
Date: Tue, 13 Aug 2024 14:26:52 +0200
Subject: [PATCH] pex-cli: add [pex-cli].args option to pass arguments to the
 PEX process globally (#21202)

A new option `[pex-cli].global_args` has been added to be able to pass
arbitrary arguments to the `pex` tool as part of any Pants goal
invocation. This should make it a lot easier to modify behavior of `pex`
tool without needing to make changes in the Pants codebase.

Not having this ability to pass arbitrary arguments to pex makes it
really hard to start taking advantage of new features that come with
newer versions of pex. For instance, the new `--exclude` flag added
recently would require making lots of changes in the codebase to be able
to pass those extra arguments. This is because the pex invocations in
the Pants codebase are numerous and it's really hard to make sure a
particular argument is respected (by keeping the chain of calls correct
making sure it does reach the final subprocess spawn). And if it's
relevant for multiple goals, this becomes even harder.

We would still need to make changes to pass arguments to individual
targets, see https://github.com/pantsbuild/pants/pull/20737 or
https://github.com/pantsbuild/pants/pull/20939 - this makes sense as
those arguments apply only to those targets. However, some options would
need to apply for any `pex` invocation (e.g. ignoring all 3rd party
dependencies).

I've looked into having environment variables support for all flags that
PEX has (https://github.com/pex-tool/pex/issues/2242) first (so that no
changes are needed in Pants, one would just export a bunch of env vars
as needed), but this is not going to happen any time soon, so doing it
in the Pants codebase instead is the only path forward, I reckon.
---
 docs/notes/2.23.x.md                          |  4 +++
 .../goals/run_pex_binary_integration_test.py  | 28 +++++++++++++++++++
 .../backend/python/util_rules/pex_cli.py      | 16 +++++++++++
 .../backend/python/util_rules/pex_cli_test.py | 10 +++++++
 4 files changed, 58 insertions(+)

diff --git a/docs/notes/2.23.x.md b/docs/notes/2.23.x.md
index cd37d8051d0..dfd8fba62c4 100644
--- a/docs/notes/2.23.x.md
+++ b/docs/notes/2.23.x.md
@@ -124,6 +124,10 @@ The default version of the `ruff` tool has been updated from 0.4.4 to 0.4.9.
 
 The default version of the [Pex](https://docs.pex-tool.org/) tool has been updated from 2.3.1 to [2.11.0](https://github.com/pex-tool/pex/releases/tag/v2.11.0).
 
+[A new option `[pex-cli].global_args`](https://www.pantsbuild.org/2.23/reference/subsystems/pex-cli#global_args) has been 
+added to be able to pass arbitrary arguments to the `pex` tool as part of any Pants goal invocation. 
+This should make it a lot easier to modify behavior of `pex` tool without needing to make changes in the Pants codebase.
+
 Fix running python source files that have dashes in them (bug introduced in 2.20). For example: `pants run path/to/some-executable.py`
 
 A new `entry_point_dependencies` field is now available for `python_tests` and `python_test` targets. This allows tests
diff --git a/src/python/pants/backend/python/goals/run_pex_binary_integration_test.py b/src/python/pants/backend/python/goals/run_pex_binary_integration_test.py
index 302184a43c5..686f0a507ef 100644
--- a/src/python/pants/backend/python/goals/run_pex_binary_integration_test.py
+++ b/src/python/pants/backend/python/goals/run_pex_binary_integration_test.py
@@ -281,3 +281,31 @@ def test_run_script_from_3rdparty_dist_issue_13747() -> None:
         result = run_pants(args)
         result.assert_success()
         assert SAY in result.stdout.strip()
+
+
+def test_pass_extra_pex_cli_subsystem_global_args() -> None:
+    """Test that extra global args passed to the pex-cli subsystem propagate to the actual pex
+    invocation."""
+    sources = {
+        "src/BUILD": dedent(
+            """\
+            python_requirement(name="cowsay", requirements=["cowsay==4.0"])
+            pex_binary(name="test", script="cowsay", dependencies=[":cowsay"])
+            """
+        ),
+    }
+    with setup_tmpdir(sources) as tmpdir:
+        args = [
+            "--backend-packages=pants.backend.python",
+            "--pex-cli-global-args='--non-existing-flag-name=some-value'",
+            f"--source-root-patterns=['/{tmpdir}/src']",
+            "run",
+            f"{tmpdir}/src:test",
+            "--",
+            "moooo",
+        ]
+        result = run_pants(args)
+        result.assert_failure()
+        stderr = result.stderr.strip()
+        assert "unrecognized arguments" in stderr
+        assert "non-existing-flag-name" in stderr
diff --git a/src/python/pants/backend/python/util_rules/pex_cli.py b/src/python/pants/backend/python/util_rules/pex_cli.py
index cace0c08e92..e7dfbab10a9 100644
--- a/src/python/pants/backend/python/util_rules/pex_cli.py
+++ b/src/python/pants/backend/python/util_rules/pex_cli.py
@@ -26,9 +26,11 @@
 from pants.engine.process import Process, ProcessCacheScope
 from pants.engine.rules import Get, collect_rules, rule
 from pants.option.global_options import GlobalOptions, ca_certs_path_to_file_content
+from pants.option.option_types import ArgsListOption
 from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 from pants.util.meta import classproperty
+from pants.util.strutil import softwrap
 
 logger = logging.getLogger(__name__)
 
@@ -42,6 +44,18 @@ class PexCli(TemplatedExternalTool):
     default_url_template = "https://github.com/pex-tool/pex/releases/download/{version}/pex"
     version_constraints = ">=2.3.0,<3.0"
 
+    # extra args to be passed to the pex tool; note that they
+    # are going to apply to all invocations of the pex tool.
+    global_args = ArgsListOption(
+        example="--check=error --no-compile",
+        extra_help=softwrap(
+            """
+            Note that these apply to all invocations of the pex tool, including building `pex_binary`
+            targets, preparing `python_test` targets to run, and generating lockfiles.
+            """
+        ),
+    )
+
     @classproperty
     def default_known_versions(cls):
         return [
@@ -123,6 +137,7 @@ async def setup_pex_cli_process(
     python_native_code: PythonNativeCodeSubsystem.EnvironmentAware,
     global_options: GlobalOptions,
     pex_subsystem: PexSubsystem,
+    pex_cli_subsystem: PexCli,
     python_setup: PythonSetup,
 ) -> Process:
     tmpdir = ".tmp"
@@ -179,6 +194,7 @@ async def setup_pex_cli_process(
         *warnings_args,
         *pip_version_args,
         *resolve_args,
+        *pex_cli_subsystem.global_args,
         # NB: This comes at the end because it may use `--` passthrough args, # which must come at
         # the end.
         *request.extra_args,
diff --git a/src/python/pants/backend/python/util_rules/pex_cli_test.py b/src/python/pants/backend/python/util_rules/pex_cli_test.py
index 20f3c099b03..57a594bc032 100644
--- a/src/python/pants/backend/python/util_rules/pex_cli_test.py
+++ b/src/python/pants/backend/python/util_rules/pex_cli_test.py
@@ -42,3 +42,13 @@ def test_custom_ca_certs(tmp_path: Path, rule_runner: RuleRunner) -> None:
     chrooted_certs_file = [f for f in files if f.path == "certsfile"]
     assert len(chrooted_certs_file) == 1
     assert b"Some fake cert" == chrooted_certs_file[0].content
+
+
+def test_pass_global_args_to_pex_cli_subsystem(tmp_path: Path, rule_runner: RuleRunner) -> None:
+    """Test that arbitrary global arguments can be passed to the pex tool process."""
+    rule_runner.set_options(["--pex-cli-global-args='--foo=bar --baz --spam=eggs'"])
+    proc = rule_runner.request(
+        Process,
+        [PexCliProcess(subcommand=(), extra_args=(), description="")],
+    )
+    assert "--foo=bar --baz --spam=eggs" in proc.argv