Support injecting args and env vars in a PEX. (#1948)

Pex now supports setting defaults for environment variables and
injecting default arguments for the entry point chosen at PEX build
time. This eliminates the need for projects like
https://github.com/kwlzn/pyuwsgi_pex and makes it easy to package a
PEXed framework app where the main entry point is fixed in the
framework and the hooks are configured to point to user supplied entry
points.

This removes the last known legitimate API use of Pex and paves the way
for Pex 3 to be cut in good conscience with a CLI-only API that does not
hinder any known use cases.

Fixes #987

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=(