Support freezing command line args and env vars into pex binaries. (p…

…antsbuild#17905)

This capability already existed in Pex, this PR plumbs it
through in Pants.

This will allow, e.g., running a single gunicorn PEX with
different args for different services, without needing
a shim file.

docs/markdown/Python/python-goals/python-package-goal.md

@@ -139,6 +139,22 @@ You must explicitly add the dependencies you'd like to the `dependencies` field.
 
 This does not work with file arguments; you must use the target address, like `./pants package helloworld:black_bin`.
 
+### Injecting command-line arguments and environment variables
+
+You can use the `inject_args` and `inject_env` fields to "freeze" command-line arguments and environment variables into the PEX file. This can save you from having to create shim files around generic binaries. For example:
+
+```python myproduct/myservice/BUILD
+python_requirement(name="gunicorn", requirements=["gunicorn==20.1.0"])
+
+pex_binary(
+  name="myservice_bin",
+  script="gunicorn",
+  args=["myproduct.myservice.wsgi:app", "--name=myservice"],
+  env={"MY_ENV_VAR=1"},
+  dependencies=[":gunicorn"],
+)
+```
+
 > 🚧 PEX files may be platform-specific
 > 
 > If your code's requirements include distributions that include native code, then the resulting PEX file will only run on the platform it was built on. 

src/python/pants/backend/python/goals/package_pex_binary.py

@@ -6,10 +6,12 @@
 from typing import Tuple
 
 from pants.backend.python.target_types import (
+    PexArgsField,
     PexBinaryDefaults,
     PexCompletePlatformsField,
     PexEmitWarningsField,
     PexEntryPointField,
+    PexEnvField,
     PexExecutionMode,
     PexExecutionModeField,
     PexIgnoreErrorsField,
@@ -46,6 +48,7 @@
 )
 from pants.engine.unions import UnionMembership, UnionRule
 from pants.util.docutil import doc_url
+from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 from pants.util.strutil import softwrap
 
@@ -60,6 +63,8 @@ class PexBinaryFieldSet(PackageFieldSet, RunFieldSet):
 
     entry_point: PexEntryPointField
     script: PexScriptField
+    args: PexArgsField
+    env: PexEnvField
 
     output_path: OutputPathField
     emit_warnings: PexEmitWarningsField
@@ -150,6 +155,8 @@ async def package_pex_binary(
             addresses=[field_set.address],
             internal_only=False,
             main=resolved_entry_point.val or field_set.script.value,
+            inject_args=field_set.args.value or [],
+            inject_env=field_set.env.value or FrozenDict[str, str](),
             platforms=PexPlatforms.create_from_platforms_field(field_set.platforms),
             complete_platforms=complete_platforms,
             output_filename=output_filename,

src/python/pants/backend/python/goals/package_pex_binary_integration_test.py

@@ -110,12 +110,22 @@ def test_warn_files_targets(rule_runner: RuleRunner, caplog) -> None:
 def test_layout(rule_runner: RuleRunner, layout: PexLayout) -> None:
     rule_runner.write_files(
         {
-            "src/py/project/app.py": "print('hello')",
+            "src/py/project/app.py": dedent(
+                """\
+                import os
+                import sys
+                print(f"FOO={os.environ.get('FOO')}")
+                print(f"BAR={os.environ.get('BAR')}")
+                print(f"ARGV={sys.argv[1:]}")
+                """
+            ),
             "src/py/project/BUILD": dedent(
                 f"""\
                 python_sources(name="lib")
                 pex_binary(
                     entry_point="app.py",
+                    args=['123', 'abc'],
+                    env={{'FOO': 'xxx', 'BAR': 'yyy'}},
                     layout="{layout.value}",
                 )
                 """
@@ -136,7 +146,14 @@ def test_layout(rule_runner: RuleRunner, layout: PexLayout) -> None:
         if PexLayout.ZIPAPP is layout
         else os.path.join(expected_pex_relpath, "__main__.py"),
     )
-    assert b"hello\n" == subprocess.run([executable], check=True, stdout=subprocess.PIPE).stdout
+    stdout = dedent(
+        """\
+        FOO=xxx
+        BAR=yyy
+        ARGV=['123', 'abc']
+        """
+    ).encode()
+    assert stdout == subprocess.run([executable], check=True, stdout=subprocess.PIPE).stdout
 
 
 @pytest.fixture

src/python/pants/backend/python/target_types.py

@@ -42,6 +42,7 @@
     AsyncFieldMixin,
     BoolField,
     Dependencies,
+    DictStringToStringField,
     DictStringToStringSequenceField,
     Field,
     IntField,
@@ -376,6 +377,26 @@ def compute_value(cls, raw_value: Optional[str], address: Address) -> Optional[C
         return ConsoleScript(value)
 
 
+class PexArgsField(StringSequenceField):
+    alias = "args"
+    help = softwrap(
+        """
+        Freeze these command-line args into the PEX. Allows you to run generic entry points
+        on specific arguments without creating a shim file.
+        """
+    )
+
+
+class PexEnvField(DictStringToStringField):
+    alias = "env"
+    help = softwrap(
+        """
+        Freeze these environment variables into the PEX. Allows you to run generic entry points
+        on a specific environment without creating a shim file.
+        """
+    )
+
+
 class PexPlatformsField(StringSequenceField):
     alias = "platforms"
     help = softwrap(
@@ -659,6 +680,8 @@ class PexBinary(Target):
         *_PEX_BINARY_COMMON_FIELDS,
         PexEntryPointField,
         PexScriptField,
+        PexArgsField,
+        PexEnvField,
         OutputPathField,
     )
     help = softwrap(

src/python/pants/backend/python/util_rules/pex.py

@@ -138,6 +138,8 @@ class PexRequest(EngineAwareParameter):
     sources: Digest | None
     additional_inputs: Digest
     main: MainSpecification | None
+    inject_args: tuple[str, ...]
+    inject_env: FrozenDict[str, str]
     additional_args: tuple[str, ...]
     pex_path: tuple[Pex, ...]
     description: str | None = dataclasses.field(compare=False)
@@ -156,6 +158,8 @@ def __init__(
         sources: Digest | None = None,
         additional_inputs: Digest | None = None,
         main: MainSpecification | None = None,
+        inject_args: Iterable[str] = (),
+        inject_env: Mapping[str, str] = FrozenDict(),
         additional_args: Iterable[str] = (),
         pex_path: Iterable[Pex] = (),
         description: str | None = None,
@@ -189,6 +193,8 @@ def __init__(
             directly in the Pex, but should be present in the environment when building the Pex.
         :param main: The main for the built Pex, equivalent to Pex's `-e` or '-c' flag. If
             left off, the Pex will open up as a REPL.
+        :param inject_args: Command line arguments to freeze in to the PEX.
+        :param inject_env: Environment variables to freeze in to the PEX.
         :param additional_args: Any additional Pex flags.
         :param pex_path: Pex files to add to the PEX_PATH.
         :param description: A human-readable description to render in the dynamic UI when building
@@ -207,6 +213,8 @@ def __init__(
         self.sources = sources
         self.additional_inputs = additional_inputs or EMPTY_DIGEST
         self.main = main
+        self.inject_args = tuple(inject_args)
+        self.inject_env = FrozenDict(inject_env)
         self.additional_args = tuple(additional_args)
         self.pex_path = tuple(pex_path)
         self.description = description
@@ -517,6 +525,11 @@ async def build_pex(
     if request.main is not None:
         argv.extend(request.main.iter_pex_args())
 
+    for injected_arg in request.inject_args:
+        argv.extend(["--inject-args", str(injected_arg)])
+    for k, v in sorted(request.inject_env.items()):
+        argv.extend(["--inject-env", f"{k}={v}"])
+
     # TODO(John Sirois): Right now any request requirements will shadow corresponding pex path
     #  requirements, which could lead to problems. Support shading python binaries.
     #  See: https://github.com/pantsbuild/pants/issues/9206

src/python/pants/backend/python/util_rules/pex_from_targets.py

@@ -6,7 +6,7 @@
 import dataclasses
 import logging
 from dataclasses import dataclass
-from typing import Iterable
+from typing import Iterable, Mapping
 
 from packaging.utils import canonicalize_name as canonicalize_project_name
 
@@ -50,6 +50,7 @@
 from pants.engine.rules import Get, MultiGet, collect_rules, rule, rule_helper
 from pants.engine.target import Target, TransitiveTargets, TransitiveTargetsRequest
 from pants.util.docutil import doc_url
+from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 from pants.util.meta import frozen_after_init
 from pants.util.strutil import path_safe, softwrap
@@ -65,6 +66,8 @@ class PexFromTargetsRequest:
     internal_only: bool
     layout: PexLayout | None
     main: MainSpecification | None
+    inject_args: tuple[str, ...]
+    inject_env: FrozenDict[str, str]
     platforms: PexPlatforms
     complete_platforms: CompletePlatforms
     additional_args: tuple[str, ...]
@@ -87,6 +90,8 @@ def __init__(
         internal_only: bool,
         layout: PexLayout | None = None,
         main: MainSpecification | None = None,
+        inject_args: Iterable[str] = (),
+        inject_env: Mapping[str, str] = FrozenDict(),
         platforms: PexPlatforms = PexPlatforms(),
         complete_platforms: CompletePlatforms = CompletePlatforms(),
         additional_args: Iterable[str] = (),
@@ -112,6 +117,8 @@ def __init__(
         :param layout: The filesystem layout to create the PEX with.
         :param main: The main for the built Pex, equivalent to Pex's `-e` or `-c` flag. If
             left off, the Pex will open up as a REPL.
+        :param inject_args: Command line arguments to freeze in to the PEX.
+        :param inject_env: Environment variables to freeze in to the PEX.
         :param platforms: Which platforms should be supported. Setting this value will cause
             interpreter constraints to not be used because platforms already constrain the valid
             Python versions, e.g. by including `cp36m` in the platform string.
@@ -139,6 +146,8 @@ def __init__(
         self.internal_only = internal_only
         self.layout = layout
         self.main = main
+        self.inject_args = tuple(inject_args)
+        self.inject_env = FrozenDict(inject_env)
         self.platforms = platforms
         self.complete_platforms = complete_platforms
         self.additional_args = tuple(additional_args)
@@ -515,6 +524,8 @@ async def create_pex_from_targets(
         platforms=request.platforms,
         complete_platforms=request.complete_platforms,
         main=request.main,
+        inject_args=request.inject_args,
+        inject_env=request.inject_env,
         sources=merged_sources_digest,
         additional_inputs=additional_inputs,
         additional_args=additional_args,