Download using PIP

Summary:
`pkg_resources.working_set.resolve()` takes an `installer` callback which is used to download a requirement if it isn't present. Add an optional (disabled with `--no-download`) installer to `bdist_xar` that uses pip to download the dependencies.
* pip is located using `pkg_resources` to make sure we get the right pip for the Python environment.
* We call the pip entry point in a separate process to download (but not install) the requirement and its dependencies.
* We add some pkg_resources finders to make `pkg_resources.find_distributions()` work with wheel archives.
* Extract and build every source distribution that we installed. We call the distributions `setup.py` with small modifications to build the wheel.
* Add each wheel to the `working_set` so we can find it later, and return the newly downloaded wheel distribution.

Since we use the `pip` installed in the running (virtual) environment, we pick up the pip config automatically. For example we automatically download from Facebook's PyPI instance when in a `pyenv`.

Reviewed By: cooperlees

Differential Revision: D8776579

fbshipit-source-id: 09c83b3081111b408757edd579ef79cb05c1d258

setup.py

@@ -25,6 +25,7 @@ def get_long_description():
     license="BSD",
     packages=["xar", "xar.commands"],
     install_requires=[
+        "pip>=10.0.1",
         # Version 34.1 fixes a bug in the dependency resolution. If this is
         # causing an problem for you, please open an issue, and we can evaluate
         # a workaround. (grep setuptools>=34.1 to see issue)

xar/commands/bdist_xar.py

@@ -9,14 +9,13 @@
 import copy
 import os
 import sys
-import zipimport
 from distutils import log
 from distutils.dir_util import mkpath, remove_tree
 from distutils.errors import DistutilsOptionError
 
 import pkg_resources
 from setuptools import Command
-from xar import py_util, xar_builder, xar_util
+from xar import finders, pip_installer, py_util, xar_builder, xar_util
 
 
 class bdist_xar(Command):
@@ -37,6 +36,7 @@ class bdist_xar(Command):
             "Default: build the script with the package name, or if there is "
             "only one console script build that, otherwise fail.",
         ),
+        ("download", None, "Download missing dependencies using pip"),
         (
             "xar-exec=",
             None,
@@ -92,6 +92,7 @@ def initialize_options(self):
         self.skip_build = False
         self.console_scripts = None
         self.interpreter = None
+        self.download = False
         # XAR options
         self.xar_exec = None
         self.xar_mount_root = None
@@ -120,17 +121,18 @@ def finalize_options(self):
             self.sqopts.zstd_level = self.xar_zstd_level
         self.xar_outputs = []
 
+        self.working_set = pkg_resources.WorkingSet(sys.path)
+        self.installer = None
+        if self.download:
+            bdist_pip = os.path.join(self.bdist_dir, "downloads")
+            mkpath(bdist_pip)
+            self.installer = pip_installer.PipInstaller(
+                bdist_pip, self.working_set, log
+            )
+
     def get_outputs(self):
         return self.xar_outputs
 
-    def _distribution_from_wheel(self, wheel):
-        importer = zipimport.zipimporter(wheel)
-        metadata = py_util.WheelMetadata(importer)
-        dist = pkg_resources.DistInfoDistribution.from_filename(
-            wheel, metadata=metadata
-        )
-        return dist
-
     def _add_distribution(self, xar):
         bdist_wheel = self.reinitialize_command("bdist_wheel")
         bdist_wheel.skip_build = self.skip_build
@@ -143,7 +145,7 @@ def _add_distribution(self, xar):
         self.run_command("bdist_wheel")
         assert len(bdist_wheel.distribution.dist_files) == 1
         wheel = bdist_wheel.distribution.dist_files[0][2]
-        dist = self._distribution_from_wheel(wheel)
+        dist = py_util.Wheel(location=wheel).distribution
         xar.add_distribution(dist)
         return dist
 
@@ -182,14 +184,21 @@ def _set_entry_point(self, xar, entry_point):
     def _deps(self, dist, extras=()):
         requires = dist.requires(extras=extras)
         try:
+            finders.register_finders()
             # Requires setuptools>=34.1 for the bug fix.
-            return set(pkg_resources.working_set.resolve(requires, extras=extras))
+            return set(
+                self.working_set.resolve(
+                    requires, extras=extras, installer=self.installer
+                )
+            )
         except pkg_resources.DistributionNotFound:
             name = self.distribution.get_name()
             requires_str = "\n\t".join(str(req) for req in requires)
             log.error(
-                "%s's requirements are not satisfied "
-                "(try 'pip install /path/to/%s'):\n\t%s" % (name, name, requires_str)
+                "%s's requirements are not satisfied:\n\t%s\n"
+                "Either pass --download to bdist_xar to download missing "
+                "dependencies with pip or try 'pip install /path/to/%s'."
+                % (name, requires_str, name)
             )
             raise
 

xar/finders.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import absolute_import, division, print_function
+
+import os
+import pkgutil
+import zipimport
+
+import pkg_resources
+from xar import py_util
+
+
+try:
+    import importlib.machinery as importlib_machinery
+
+    # access attribute to force import under delayed import mechanisms.
+    importlib_machinery.__name__
+except ImportError:
+    importlib_machinery = None
+
+
+def find_wheels_in_zip(importer, path_item, only=False):
+    try:
+        yield py_util.Wheel(location=path_item, importer=importer).distribution
+    except Exception:
+        pass
+
+
+def find_wheels_on_path(importer, path_item, only=False):
+    if only or not os.path.isdir(path_item) or not os.access(path_item, os.R_OK):
+        return
+    for entry in os.listdir(path_item):
+        if py_util.Wheel.is_wheel_archive(entry):
+            location = os.path.join(path_item, entry)
+            for dist in pkg_resources.find_distributions(location):
+                yield dist
+
+
+def find_on_path(importer, path_item, only=False):
+    for finder in (pkg_resources.find_on_path, find_wheels_on_path):
+        for dist in finder(importer, path_item, only):
+            yield dist
+
+
+__REGISTERED = False
+
+
+def register_finders():
+    """
+    Register pkg_resources finders that work with wheels (but not eggs). These
+    replace the default pkg_resources finders. This function should be called
+    before calling pkg_resources.find_distributions().
+    """
+    global __REGISTERED
+    if __REGISTERED:
+        return
+
+    pkg_resources.register_finder(zipimport.zipimporter, find_wheels_in_zip)
+    pkg_resources.register_finder(pkgutil.ImpImporter, find_on_path)
+    if hasattr(importlib_machinery, "FileFinder"):
+        pkg_resources.register_finder(importlib_machinery.FileFinder, find_on_path)
+
+    __REGISTERED = True

xar/pip_installer.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import absolute_import, division, print_function
+
+import contextlib
+import multiprocessing
+import os
+import subprocess
+import sys
+import tarfile
+import tempfile
+import zipfile
+
+import pkg_resources
+from xar import finders, py_util, xar_util
+
+
+class PipException(Exception):
+    pass
+
+
+class BuildException(Exception):
+    pass
+
+
+class PipInstaller(object):
+    """
+    Installer function object for pkg_resources.working_set.resolve().
+    It is called like `installer(requirement)` when the requirement can't be
+    found in the working set. See `__call__()` for documentation.
+    """
+
+    def __init__(self, dest, working_set, log=None):
+        """
+        Download Wheels to `dest` and add them to the `working_set`.
+        """
+        self._dest = dest
+        self._working_set = working_set
+        self._log = log
+        self._working_set.add_entry(self._dest)
+        req = pkg_resources.Requirement.parse("pip")
+        dist = self._working_set.find(req)
+        self._pip_main = pkg_resources.load_entry_point(dist, "console_scripts", "pip")
+
+    def clean(self):
+        """
+        Remove any non-wheels from the downloads directory.
+        """
+        for entry in os.listdir(self._dest):
+            if not py_util.Wheel.is_wheel_archive(entry):
+                xar_util.safe_remove(os.path.join(self._dest, entry))
+
+    def invoke_pip(self, args):
+        def main(args):
+            sys.exit(self._pip_main(args))
+
+        p = multiprocessing.Process(target=main, args=(args,))
+        p.start()
+        p.join()
+        if p.exitcode == 0:
+            return
+        raise PipException("'pip %s' failed" % " ".join(args))
+
+    def download(self, requirement):
+        """
+        Download the requirement to the downloads directory using pip.
+        """
+        args = ["download", "-d", self._dest, str(requirement)]
+        self.invoke_pip(args)
+
+    def extract_sdist(self, sdist, dest):
+        """
+        Extract the sdist archive and return the path to the source.
+        """
+        if sdist.lower().endswith(".zip"):
+            open_sdist = zipfile.ZipFile
+            error_cls = zipfile.BadZipfile
+        else:
+            assert ".tar" in sdist.lower()
+            open_sdist = tarfile.TarFile.open
+            error_cls = tarfile.ReadError
+        try:
+            with contextlib.closing(open_sdist(sdist)) as archive:
+                archive.extractall(path=dest)
+
+            def collapse_trivial(path):
+                entries = os.listdir(path)
+                if len(entries) == 1:
+                    entry = os.path.join(path, entries[0])
+                    if os.path.isdir(entry):
+                        return collapse_trivial(entry)
+                return path
+
+            return collapse_trivial(dest)
+        except error_cls:
+            raise BuildException("Failed to extract %s" % os.path.basename(sdist))
+
+    def build_wheel_from_sdist(self, sdist):
+        """
+        Given a sdist archive extract it, build in a temporary directory, and
+        put the wheel into the downloads directory.
+        """
+        temp = tempfile.mkdtemp()
+        try:
+            source = self.extract_sdist(sdist, temp)
+            # Make sure to import setuptools and wheel in the setup.py.
+            # This is happening in a temporary directory, so we will just
+            # overwrite the setup.py to add our own imports.
+            setup_py = os.path.join(source, "setup.py")
+            with open(setup_py, "r") as f:
+                original = f.read()
+            with open(setup_py, "w") as f:
+                f.write("import setuptools\n")
+                f.write("import wheel\n")
+                f.write(original)
+            # Build the wheel
+            command = [
+                sys.executable,
+                setup_py,
+                "bdist_wheel",
+                "-d",
+                os.path.abspath(self._dest),
+            ]
+            subprocess.check_call(command, cwd=source)
+        except subprocess.CalledProcessError:
+            raise BuildException("Failed to build %s" % str(os.path.basename(sdist)))
+        finally:
+            xar_util.safe_rmtree(temp)
+
+    def find(self, requirement):
+        """
+        Ensure all built wheels are added to the working set.
+        Return the distribution.
+        """
+        finders.register_finders()
+        for dist in pkg_resources.find_distributions(self._dest):
+            if dist not in self._working_set:
+                self._working_set.add(dist, entry=self._dest)
+        return self._working_set.find(requirement)
+
+    def __call__(self, requirement):
+        """
+        Attempts to download the requirement (and its dependencies) and add the
+        wheel(s) to the downloads directory and the working set. Returns the
+        distribution on success and None on failure.
+        """
+        # Remove non-wheels from the download directory
+        self.clean()
+        # Attempt to download the wheel/sdist
+        try:
+            self.download(requirement)
+        except PipException as e:
+            if self._log:
+                self._log.exception(e)
+            return None
+        # Build wheels for the sdists (and remove the sdist)
+        for entry in os.listdir(self._dest):
+            if py_util.Wheel.is_wheel_archive(entry):
+                continue
+            try:
+                sdist = os.path.join(self._dest, entry)
+                self.build_wheel_from_sdist(sdist)
+                xar_util.safe_remove(sdist)
+            except BuildException as e:
+                if self._log:
+                    self._log.exception(e)
+                return None
+        # Return the wheel distribution
+        return self.find(requirement)