Add support for PEX runtime tools & an info tool. (pex-tool#1127)

Add a new `--include-tools` option to include any pex.tools in generated
PEX files. These tools are activated by running PEX files with
PEX_TOOLS=1. The `Info` tool seeds the tool set and simply dumps the
effective PEX-INFO for the given PEX.

Work towards pex-tool#962 and pex-tool#1115

pex/bin/pex.py

@@ -295,6 +295,15 @@ def configure_clp_pex_options(parser):
         "Tailor the behavior of the emitted .pex file if -o is specified.",
     )
 
+    group.add_argument(
+        "--include-tools",
+        dest="include_tools",
+        default=False,
+        action=HandleBoolAction,
+        help="Whether to include runtime tools in the pex file. If included, these can be run by "
+        "exporting PEX_TOOLS=1 and following the usage and --help information.",
+    )
+
     group.add_argument(
         "--zip-safe",
         "--not-zip-safe",
@@ -775,7 +784,12 @@ def to_python_interpreter(full_path_or_basename):
         # options.preamble_file is None
         preamble = None
 
-    pex_builder = PEXBuilder(path=safe_mkdtemp(), interpreter=interpreter, preamble=preamble)
+    pex_builder = PEXBuilder(
+        path=safe_mkdtemp(),
+        interpreter=interpreter,
+        preamble=preamble,
+        include_tools=options.include_tools,
+    )
 
     if options.resources_directory:
         pex_warnings.warn(

pex/pex.py

@@ -27,7 +27,7 @@
 from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
 from pex.util import iter_pth_paths, named_temporary_file
-from pex.variables import ENV
+from pex.variables import ENV, Variables
 
 if TYPE_CHECKING:
     from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Set, Tuple, TypeVar
@@ -65,25 +65,37 @@ def _clean_environment(cls, env=None, strip_pex_env=True):
                 if key.startswith("PEX_"):
                     del env[key]
 
-    def __init__(self, pex=sys.argv[0], interpreter=None, env=ENV, verify_entry_point=False):
+    def __init__(
+        self,
+        pex=sys.argv[0],  # type: str
+        interpreter=None,  # type: Optional[PythonInterpreter]
+        env=ENV,  # type: Variables
+        verify_entry_point=False,  # type: bool
+    ):
+        # type: (...) -> None
         self._pex = pex
         self._interpreter = interpreter or PythonInterpreter.get()
         self._pex_info = PexInfo.from_pex(self._pex)
         self._pex_info_overrides = PexInfo.from_env(env=env)
         self._vars = env
-        self._envs = []
-        self._working_set = None
+        self._envs = []  # type: List[PEXEnvironment]
+        self._working_set = None  # type: Optional[WorkingSet]
         if verify_entry_point:
             self._do_entry_point_verification()
 
+    def pex_info(self):
+        # type: () -> PexInfo
+        pex_info = self._pex_info.copy()
+        pex_info.update(self._pex_info_overrides)
+        pex_info.merge_pex_path(self._vars.PEX_PATH)
+        return pex_info
+
     def _activate(self):
         if not self._working_set:
             working_set = WorkingSet([])
 
             # set up the local .pex environment
-            pex_info = self._pex_info.copy()
-            pex_info.update(self._pex_info_overrides)
-            pex_info.merge_pex_path(self._vars.PEX_PATH)
+            pex_info = self.pex_info()
             self._envs.append(PEXEnvironment(self._pex, pex_info, interpreter=self._interpreter))
             # N.B. by this point, `pex_info.pex_path` will contain a single pex path
             # merged from pex_path in `PEX-INFO` and `PEX_PATH` set in the environment.
@@ -426,6 +438,7 @@ def _wrap_profiling(self, runner, *args):
                 profiler.print_stats(sort=pex_profile_sort)
 
     def path(self):
+        # type: () -> str
         """Return the path this PEX was built at."""
         return self._pex
 
@@ -436,19 +449,37 @@ def execute(self):
         This function makes assumptions that it is the last function called by the interpreter.
         """
         teardown_verbosity = self._vars.PEX_TEARDOWN_VERBOSE
+
+        # N.B.: This is set in `__main__.py` of the executed PEX by `PEXBuilder` when we've been
+        # executed from within a PEX zip file in `--unzip` mode.  We replace `sys.argv[0]` to avoid
+        # confusion and allow the user code we hand off to to provide useful messages and fully
+        # valid re-execs that always re-directed through the PEX file.
+        sys.argv[0] = os.environ.pop("__PEX_EXE__", sys.argv[0])
+
         try:
             pex_inherit_path = self._vars.PEX_INHERIT_PATH
             if pex_inherit_path == InheritPath.FALSE:
                 pex_inherit_path = self._pex_info.inherit_path
             self.patch_sys(pex_inherit_path)
             working_set = self._activate()
             self.patch_pkg_resources(working_set)
-            exit_code = self._wrap_coverage(self._wrap_profiling, self._execute)
+            if self._vars.PEX_TOOLS:
+                try:
+                    from pex.tools import main as tools
+                except ImportError as e:
+                    die(
+                        "This PEX was not built with tools (Re-build the PEX file with "
+                        "`pex --include-tools ...`): {}".format(e)
+                    )
+
+                exit_code = tools.main(pex=self, pex_prog_path=sys.argv[0])
+            else:
+                exit_code = self._wrap_coverage(self._wrap_profiling, self._execute)
             if exit_code:
                 sys.exit(exit_code)
         except Exception:
-            # Allow the current sys.excepthook to handle this app exception before we tear things down in
-            # finally, then reraise so that the exit status is reflected correctly.
+            # Allow the current sys.excepthook to handle this app exception before we tear things
+            # down in finally, then reraise so that the exit status is reflected correctly.
             sys.excepthook(*sys.exc_info())
             raise
         except SystemExit as se:
@@ -480,12 +511,6 @@ def execute(self):
     def _execute(self):
         force_interpreter = self._vars.PEX_INTERPRETER
 
-        # N.B.: This is set in `__main__.py` of the executed PEX by `PEXBuilder` when we've been
-        # executed from within a PEX zip file in `--unzip` mode.  We replace `sys.argv[0]` to avoid
-        # confusion and allow the user code we hand off to to provide useful messages and fully valid
-        # re-execs that always re-directed through the PEX file.
-        sys.argv[0] = os.environ.pop("__PEX_EXE__", sys.argv[0])
-
         self._clean_environment(strip_pex_env=self._pex_info.strip_pex_env)
 
         if force_interpreter:

pex/pex_builder.py

@@ -116,7 +116,14 @@ class InvalidExecutableSpecification(Error):
         pass
 
     def __init__(
-        self, path=None, interpreter=None, chroot=None, pex_info=None, preamble=None, copy=False
+        self,
+        path=None,
+        interpreter=None,
+        chroot=None,
+        pex_info=None,
+        preamble=None,
+        copy=False,
+        include_tools=False,
     ):
         """Initialize a pex builder.
 
@@ -131,6 +138,8 @@ def __init__(
         :type preamble: str
         :keyword copy: If False, attempt to create the pex environment via hard-linking, falling
                        back to copying across devices. If True, always copy.
+        :keyword include_tools: If True, include runtime tools which can be executed by exporting
+                                `PEX_TOOLS=1`.
 
         .. versionchanged:: 0.8
           The temporary directory created when ``path`` is not specified is now garbage collected on
@@ -141,6 +150,7 @@ def __init__(
         self._pex_info = pex_info or PexInfo.default(self._interpreter)
         self._preamble = preamble or ""
         self._copy = copy
+        self._include_tools = include_tools
 
         self._shebang = self._interpreter.identity.hashbang()
         self._logger = logging.getLogger(__name__)
@@ -445,9 +455,7 @@ def _precompile_source(self):
             self._chroot.touch(compiled, label="bytecode")
 
     def _prepare_manifest(self):
-        self._chroot.write(
-            self._pex_info.dump(sort_keys=True).encode("utf-8"), PexInfo.PATH, label="manifest"
-        )
+        self._chroot.write(self._pex_info.dump().encode("utf-8"), PexInfo.PATH, label="manifest")
 
     def _prepare_main(self):
         self._chroot.write(
@@ -480,9 +488,13 @@ def _prepare_bootstrap(self):
         if not isinstance(provider, DefaultProvider):
             mod = __import__(source_name, fromlist=["ignore"])
             provider = ZipProvider(mod)
-        for package in ("", "third_party"):
+
+        bootstrap_packages = ["", "third_party"]
+        if self._include_tools:
+            bootstrap_packages.extend(["tools", "tools/commands"])
+        for package in bootstrap_packages:
             for fn in provider.resource_listdir(package):
-                if fn.endswith(".py"):
+                if not (provider.resource_isdir(os.path.join(package, fn)) or fn.endswith(".pyc")):
                     rel_path = os.path.join(package, fn)
                     self._chroot.write(
                         provider.get_resource_string(source_name, rel_path),

pex/pex_info.py

@@ -12,14 +12,14 @@
 from pex.compatibility import string as compatibility_string
 from pex.inherit_path import InheritPath
 from pex.orderedset import OrderedSet
-from pex.typing import TYPE_CHECKING
+from pex.typing import TYPE_CHECKING, cast
 from pex.variables import ENV, Variables
 from pex.version import __version__ as pex_version
 
 if TYPE_CHECKING:
     from pex.interpreter import PythonInterpreter
 
-    from typing import Any, Mapping, Optional, Text, Union
+    from typing import Any, Dict, Mapping, Optional, Text, Union
 
 
 # TODO(wickman) Split this into a PexInfoBuilder/PexInfo to ensure immutability.
@@ -89,6 +89,7 @@ def default(cls, interpreter=None):
 
     @classmethod
     def from_pex(cls, pex):
+        # type: (str) -> PexInfo
         if os.path.isfile(pex):
             with open_zip(pex) as zf:
                 pex_info = zf.read(cls.PATH)
@@ -151,7 +152,7 @@ def __init__(self, info=None):
             raise ValueError(
                 "PexInfo can only be seeded with a dict, got: " "%s of type %s" % (info, type(info))
             )
-        self._pex_info = dict(info) if info else {}  # type Dict[str, str]
+        self._pex_info = dict(info) if info else {}  # type Dict[str, Any]
         self._distributions = self._pex_info.get("distributions", {})
         # cast as set because pex info from json must store interpreter_constraints as a list
         self._interpreter_constraints = set(self._pex_info.get("interpreter_constraints", set()))
@@ -226,15 +227,17 @@ def strip_pex_env(self, value):
 
     @property
     def pex_path(self):
+        # type: () -> Optional[str]
         """A colon separated list of other pex files to merge into the runtime environment.
 
         This pex info property is used to persist the PEX_PATH environment variable into the pex
         info metadata for reuse within a built pex.
         """
-        return self._pex_info.get("pex_path")
+        return cast("Optional[str]", self._pex_info.get("pex_path"))
 
     @pex_path.setter
     def pex_path(self, value):
+        # type: (str) -> None
         self._pex_info["pex_path"] = value
 
     @property
@@ -365,24 +368,33 @@ def zip_unsafe_cache(self):
         return os.path.join(self.pex_root, "code")
 
     def update(self, other):
+        # type: (PexInfo) -> None
         if not isinstance(other, PexInfo):
             raise TypeError("Cannot merge a %r with PexInfo" % type(other))
         self._pex_info.update(other._pex_info)
         self._distributions.update(other.distributions)
         self._interpreter_constraints.update(other.interpreter_constraints)
         self._requirements.update(other.requirements)
 
-    def dump(self, sort_keys=False):
-        # type: (bool) -> str
-        pex_info_copy = self._pex_info.copy()
-        pex_info_copy["inherit_path"] = self.inherit_path.value
-        pex_info_copy["requirements"] = sorted(self._requirements)
-        pex_info_copy["interpreter_constraints"] = sorted(self._interpreter_constraints)
-        pex_info_copy["distributions"] = self._distributions.copy()
-        return json.dumps(pex_info_copy, sort_keys=sort_keys)
+    def as_json_dict(self):
+        # type: () -> Dict[str, Any]
+        data = self._pex_info.copy()
+        data["inherit_path"] = self.inherit_path.value
+        data["requirements"] = list(self._requirements)
+        data["interpreter_constraints"] = list(self._interpreter_constraints)
+        data["distributions"] = self._distributions.copy()
+        return data
+
+    def dump(self):
+        # type: (...) -> str
+        data = self.as_json_dict()
+        data["requirements"].sort()
+        data["interpreter_constraints"].sort()
+        return json.dumps(data, sort_keys=True)
 
     def copy(self):
-        return self.from_json(self.dump())
+        # type: () -> PexInfo
+        return PexInfo(self._pex_info)
 
     @staticmethod
     def _merge_split(*paths):

pex/tools/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).

pex/tools/__main__.py

@@ -0,0 +1,10 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+import sys
+
+from pex.tools.main import main as tools
+
+sys.exit(tools())