Support injecting args and env vars in a PEX.

docs/recipes.rst

@@ -3,6 +3,40 @@
 PEX Recipes and Notes
 =====================
 
+Uvicorn and other customizable application servers
+--------------------------------------------------
+
+Often you want to run a third-party application server and have it use your code. You can always do
+this by writing a shim bit of python code that starts the application server configured to use your
+code. It may be simpler though to use ``--inject-env`` and ``--inject-args`` to seal this
+configuration into a PEX file without needing to write a shim.
+
+For example, to package up a uvicorn-powered server of your app coroutine in ``example.py`` that ran
+on port 8888 by default you could:
+
+.. code-block:: bash
+
+    $ pex "uvicorn[standard]" -c uvicorn --inject-args 'example:app --port 8888' -oexample-app.pex
+    $ ./example-app.pex
+    INFO:     Started server process [2014]
+    INFO:     Waiting for application startup.
+    INFO:     ASGI 'lifespan' protocol appears unsupported.
+    INFO:     Application startup complete.
+    INFO:     Uvicorn running on http://127.0.0.1:8888 (Press CTRL+C to quit)
+    ^CINFO:     Shutting down
+    INFO:     Finished server process [2014]
+
+You could then over-ride the port with:
+
+.. code-block:: bash
+
+    $ ./example-app.pex --port 0
+    INFO:     Started server process [2248]
+    INFO:     Waiting for application startup.
+    INFO:     ASGI 'lifespan' protocol appears unsupported.
+    INFO:     Application startup complete.
+    INFO:     Uvicorn running on http://127.0.0.1:45751 (Press CTRL+C to quit)
+
 Long running PEX applications and daemons
 -----------------------------------------
 

pex/bin/pex.py

@@ -11,8 +11,9 @@
 import itertools
 import json
 import os
+import shlex
 import sys
-from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser
+from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentError, ArgumentParser
 from textwrap import TextWrapper
 
 from pex import pex_warnings
@@ -45,7 +46,6 @@
 from pex.resolve.resolvers import Unsatisfiable
 from pex.resolver import resolve
 from pex.result import catch, try_
-from pex.sh_boot import create_sh_boot_script
 from pex.targets import Targets
 from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING, cast
@@ -405,6 +405,37 @@ def configure_clp_pex_entry_points(parser):
         "`from a.b.c import m` during validation.",
     )
 
+    class InjectEnvAction(Action):
+        def __call__(self, parser, namespace, value, option_str=None):
+            components = value.split("=", 1)
+            if len(components) != 2:
+                raise ArgumentError(
+                    self,
+                    "Environment variable values must be of the form `name=value`. "
+                    "Given: {value}".format(value=value),
+                )
+            self.default.append(tuple(components))
+
+    group.add_argument(
+        "--inject-env",
+        dest="inject_env",
+        default=[],
+        action=InjectEnvAction,
+        help="Environment variables to freeze in to the application environment.",
+    )
+
+    class InjectArgAction(Action):
+        def __call__(self, parser, namespace, value, option_str=None):
+            self.default.extend(shlex.split(value))
+
+    group.add_argument(
+        "--inject-args",
+        dest="inject_args",
+        default=[],
+        action=InjectArgAction,
+        help="Command line arguments to the application to freeze in.",
+    )
+
 
 class Seed(Enum["Seed.Value"]):
     class Value(Enum.Value):
@@ -603,6 +634,8 @@ def build_pex(
                 pex_builder.add_source(src_file_path, dst_path)
 
     pex_info = pex_builder.info
+    pex_info.inject_env = dict(options.inject_env)
+    pex_info.inject_args = options.inject_args
     pex_info.venv = bool(options.venv)
     pex_info.venv_bin_path = options.venv or BinPath.FALSE
     pex_info.venv_copies = options.venv_copies

pex/pex.py

@@ -560,6 +560,17 @@ def _execute(self):
             TRACER.log("PEX_INTERPRETER specified, dropping into interpreter")
             return self.execute_interpreter()
 
+        if not any(
+            (
+                self._pex_info_overrides.script,
+                self._pex_info_overrides.entry_point,
+                self._pex_info.script,
+                self._pex_info.entry_point,
+            )
+        ):
+            TRACER.log("No entry point specified, dropping into interpreter")
+            return self.execute_interpreter()
+
         if self._pex_info_overrides.script and self._pex_info_overrides.entry_point:
             return "Cannot specify both script and entry_point for a PEX!"
 
@@ -568,19 +579,21 @@ def _execute(self):
 
         if self._pex_info_overrides.script:
             return self.execute_script(self._pex_info_overrides.script)
-        elif self._pex_info_overrides.entry_point:
+        if self._pex_info_overrides.entry_point:
             return self.execute_entry(
                 EntryPoint.parse("run = {}".format(self._pex_info_overrides.entry_point))
             )
-        elif self._pex_info.script:
+
+        for name, value in self._pex_info.inject_env.items():
+            os.environ.setdefault(name, value)
+        sys.argv[1:1] = list(self._pex_info.inject_args)
+
+        if self._pex_info.script:
             return self.execute_script(self._pex_info.script)
-        elif self._pex_info.entry_point:
+        else:
             return self.execute_entry(
                 EntryPoint.parse("run = {}".format(self._pex_info.entry_point))
             )
-        else:
-            TRACER.log("No entry point specified, dropping into interpreter")
-            return self.execute_interpreter()
 
     @classmethod
     def demote_bootstrap(cls):

pex/pex_info.py

@@ -165,6 +165,26 @@ def build_properties(self, value):
         self._pex_info["build_properties"] = self.make_build_properties()
         self._pex_info["build_properties"].update(value)
 
+    @property
+    def inject_env(self):
+        # type: () -> Dict[str, str]
+        return dict(self._pex_info.get("inject_env", {}))
+
+    @inject_env.setter
+    def inject_env(self, value):
+        # type: (Mapping[str, str]) -> None
+        self._pex_info["inject_env"] = dict(value)
+
+    @property
+    def inject_args(self):
+        # type: () -> Tuple[str, ...]
+        return tuple(self._pex_info.get("inject_args", ()))
+
+    @inject_args.setter
+    def inject_args(self, value):
+        # type: (Iterable[str]) -> None
+        self._pex_info["inject_args"] = tuple(value)
+
     @property
     def venv(self):
         # type: () -> bool

pex/venv/pex.py

@@ -470,6 +470,8 @@ def sys_executable_paths():
                     )
                 )
                 sys.exit(1)
+            is_exec_override = len(pex_overrides) == 1
+
             if {strip_pex_env!r}:
                 for key in list(os.environ):
                     if key.startswith("PEX_"):
@@ -559,6 +561,11 @@ def sys_executable_paths():
                     {exec_ast}
                     sys.exit(0)
 
+            if not is_exec_override:
+                for name, value in {inject_env!r}:
+                    os.environ.setdefault(name, value)
+                sys.argv[1:1] = {inject_args!r}
+
             module_name, _, function = entry_point.partition(":")
             if not function:
                 import runpy
@@ -578,6 +585,8 @@ def sys_executable_paths():
             shebang_python=venv_python,
             bin_path=bin_path,
             strip_pex_env=pex_info.strip_pex_env,
+            inject_env=tuple(pex_info.inject_env.items()),
+            inject_args=list(pex_info.inject_args),
             entry_point=pex_info.entry_point,
             script=pex_info.script,
             exec_ast=(