From 6f495a8a52010f2c317d8db4356c9b43b432c5a4 Mon Sep 17 00:00:00 2001
From: Huon Wilson <huon@exoflare.io>
Date: Tue, 2 Apr 2024 15:53:47 +1100
Subject: [PATCH] Allow passing arguments to PEX invocation used for
 pex_binaries

---
 .../python/goals/package_pex_binary.py        |  4 +++
 .../package_pex_binary_integration_test.py    | 36 ++++++++++++++++++-
 .../pants/backend/python/target_types.py      | 26 ++++++++++++--
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/src/python/pants/backend/python/goals/package_pex_binary.py b/src/python/pants/backend/python/goals/package_pex_binary.py
index 23e515a8a0a..22204129616 100644
--- a/src/python/pants/backend/python/goals/package_pex_binary.py
+++ b/src/python/pants/backend/python/goals/package_pex_binary.py
@@ -16,6 +16,7 @@
     PexExecutableField,
     PexExecutionMode,
     PexExecutionModeField,
+    PexExtraBuildArgsField,
     PexIgnoreErrorsField,
     PexIncludeRequirementsField,
     PexIncludeSourcesField,
@@ -83,6 +84,7 @@ class PexBinaryFieldSet(PackageFieldSet, RunFieldSet):
     venv_hermetic_scripts: PexVenvHermeticScripts
     environment: EnvironmentField
     check: PexCheckField
+    extra_build_args: PexExtraBuildArgsField
 
     @property
     def _execution_mode(self) -> PexExecutionMode:
@@ -116,6 +118,8 @@ def generate_additional_args(self, pex_binary_defaults: PexBinaryDefaults) -> Tu
             args.append("--venv-site-packages-copies")
         if self.venv_hermetic_scripts.value is False:
             args.append("--non-hermetic-venv-scripts")
+        if self.extra_build_args.value:
+            args.extend(self.extra_build_args.value)
         return tuple(args)
 
 
diff --git a/src/python/pants/backend/python/goals/package_pex_binary_integration_test.py b/src/python/pants/backend/python/goals/package_pex_binary_integration_test.py
index 246a95bccaf..ec4ddb6e332 100644
--- a/src/python/pants/backend/python/goals/package_pex_binary_integration_test.py
+++ b/src/python/pants/backend/python/goals/package_pex_binary_integration_test.py
@@ -14,7 +14,10 @@
 
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import package_dists, package_pex_binary
-from pants.backend.python.goals.package_pex_binary import PexBinaryFieldSet
+from pants.backend.python.goals.package_pex_binary import (
+    PexBinaryFieldSet,
+    PexFromTargetsRequestForBuiltPackage,
+)
 from pants.backend.python.macros.python_artifact import PythonArtifact
 from pants.backend.python.subsystems.setuptools import PythonDistributionFieldSet
 from pants.backend.python.target_types import (
@@ -49,6 +52,7 @@ def rule_runner() -> PythonRuleRunner:
             *core_target_types_rules(),
             *package_dists.rules(),
             QueryRule(BuiltPackage, [PexBinaryFieldSet]),
+            QueryRule(PexFromTargetsRequestForBuiltPackage, [PexBinaryFieldSet]),
             QueryRule(BuiltPackage, [PythonDistributionFieldSet]),
         ],
         target_types=[
@@ -455,3 +459,33 @@ def test_sh_boot_plumb(rule_runner: PythonRuleRunner) -> None:
     with open(executable, "rb") as f:
         shebang = f.readline().decode()
         assert "#!/bin/sh" in shebang
+
+
+def test_extra_build_args(rule_runner: PythonRuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/py/project/app.py": dedent(
+                """\
+                print("hello")
+                """
+            ),
+            "src/py/project/BUILD": dedent(
+                """\
+                python_sources(name="lib")
+                pex_binary(
+                    entry_point="app.py",
+                    extra_build_args=["--example-extra-arg", "value-goes-here"]
+                )
+                """
+            ),
+        }
+    )
+
+    tgt = rule_runner.get_target(Address("src/py/project"))
+    field_set = PexBinaryFieldSet.create(tgt)
+    result = rule_runner.request(PexFromTargetsRequestForBuiltPackage, [field_set])
+
+    additional_args = result.request.additional_args
+
+    assert additional_args[-2] == "--example-extra-arg"
+    assert additional_args[-1] == "value-goes-here"
diff --git a/src/python/pants/backend/python/target_types.py b/src/python/pants/backend/python/target_types.py
index e32e338d886..9f08a0be1b9 100644
--- a/src/python/pants/backend/python/target_types.py
+++ b/src/python/pants/backend/python/target_types.py
@@ -415,11 +415,32 @@ def compute_value(cls, raw_value: Optional[str], address: Address) -> Optional[E
 
 
 class PexArgsField(StringSequenceField):
-    alias = "args"
+    alias: ClassVar[str] = "args"
     help = help_text(
-        """
+        lambda: f"""
         Freeze these command-line args into the PEX. Allows you to run generic entry points
         on specific arguments without creating a shim file.
+
+        This is different to `{PexExtraBuildArgsField.alias}`: `{PexArgsField.alias}`
+        records arguments used by the packaged PEX when executed,
+        `{PexExtraBuildArgsField.alias}` passes arguments to the process that does the
+        packaging.
+        """
+    )
+
+
+class PexExtraBuildArgsField(StringSequenceField):
+    alias: ClassVar[str] = "extra_build_args"
+    default = ()
+    help = help_text(
+        lambda: f"""
+        Extra arguments to pass to the `pex` invocation used to build this PEX. These are
+        passed after all other arguments. This can be used to pass extra options that
+        Pants doesn't have built-in support for.
+
+        This is different to `{PexArgsField.alias}`: `{PexArgsField.alias}` records
+        arguments used by the packaged PEX when executed, `{PexExtraBuildArgsField.alias}`
+        passes arguments to the process that does the packaging.
         """
     )
 
@@ -789,6 +810,7 @@ class PexBinary(Target):
         PexScriptField,
         PexExecutableField,
         PexArgsField,
+        PexExtraBuildArgsField,
         PexEnvField,
         OutputPathField,
     )