| # Copyright 2016 Google Inc. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """Runtime support code for executables created by Subpar. |
| |
| 1. Third-party modules require some PYTHONPATH manipulation. |
| |
| 2. Python can natively import python modules from a zip archive, but |
| C extension modules require some help. |
| |
| 3. Resources stored in a .par file may need to be exposed as OS-level |
| files instead of Python File objects. |
| |
| We hook into the pkg_resources module, if present, to achieve 2 and 3. |
| |
| Limitations: |
| |
| A. Retrieving resources from packages |
| |
| It should be possible to do this: |
| fn = pkg_resources.resource_filename('mypackage', 'myfile') |
| But instead one must do |
| fn = pkg_resources.resource_filename( |
| pkg_resources.Requirement.parse.spec('mypackage'), |
| 'myfile') |
| |
| B. Extraction dir |
| |
| You should explicitly set the default extraction directory, via |
| `pkg_resources.set_extraction_path(my_directory)`, since the default |
| is not safe. For example: |
| |
| tmpdir = tempfile.mkdtemp() |
| pkg_resources.set_extraction(tmpdir) |
| |
| You should arrange for that directory to be deleted at some point. |
| Note that pkg_resources.cleanup_resources() is an unimplemented no-op, |
| so use something else. For example: |
| |
| atexit.register(lambda: shutil.rmtree(tmpdir, ignore_errors=True)) |
| |
| """ |
| |
| import atexit |
| import os |
| import pkgutil |
| import shutil |
| import sys |
| import tempfile |
| import warnings |
| import zipfile |
| import zipimport |
| |
| |
| def _log(msg): |
| """Print a debugging message in the same format as python -vv output""" |
| if sys.flags.verbose: |
| sys.stderr.write(msg) |
| sys.stderr.write('\n') |
| |
| |
| def _find_archive(): |
| """Find the path to the currently executing .par file |
| |
| We don't handle the case where prefix is non-empty. |
| """ |
| main = sys.modules.get('__main__') |
| if not main: |
| _log('# __main__ module not found') |
| return None |
| main_loader = getattr(main, '__loader__') |
| if not main_loader: |
| _log('# __main__.__loader__ not set') |
| return None |
| prefix = getattr(main_loader, 'prefix') |
| if prefix != '': |
| _log('# unexpected prefix for __main__.__loader__ is %s' % |
| main_loader.prefix) |
| return None |
| archive_path = getattr(main_loader, 'archive') |
| if not archive_path: |
| _log('# missing archive for __main__.__loader__') |
| return None |
| return archive_path |
| |
| |
| def _extract_files(archive_path): |
| """Extract the contents of this .par file to disk. |
| |
| This creates a temporary directory, and registers an atexit |
| handler to clean that directory on program exit. Extraction and |
| cleanup will potentially use significant time and disk space. |
| |
| Returns: |
| Directory where contents were extracted to. |
| """ |
| extract_dir = tempfile.mkdtemp() |
| |
| def _extract_files_cleanup(): |
| shutil.rmtree(extract_dir, ignore_errors=True) |
| atexit.register(_extract_files_cleanup) |
| _log('# extracting %s to %s' % (archive_path, extract_dir)) |
| |
| zip_file = zipfile.ZipFile(archive_path, mode='r') |
| zip_file.extractall(extract_dir) |
| zip_file.close() |
| |
| return extract_dir |
| |
| |
| def _version_check_pkg_resources(pkg_resources): |
| """Check that pkg_resources supports the APIs we need.""" |
| # Check that pkg_resources is new enough. |
| # |
| # Determining the version of an arbitrarily old version of |
| # pkg_resources is tough, since it doesn't have a version literal, |
| # and the accompanying setuptools package computes its version |
| # dynamically from metadata that might not exist. Also setuptools |
| # might not exist, especially in the case of the pip-vendored copy |
| # of pkg_resources. |
| # |
| # We do a feature detection instead. We examine |
| # pkg_resources.WorkingSet.add, and see if it has at least the |
| # third default argument ('replace'). |
| try: |
| if sys.version_info[0] < 3: |
| defaults = pkg_resources.WorkingSet.add.im_func.func_defaults |
| else: |
| defaults = pkg_resources.WorkingSet.add.__defaults__ |
| return len(defaults) >= 3 |
| except AttributeError: |
| return False |
| |
| |
| def _setup_pkg_resources(pkg_resources_name): |
| """Setup hooks into the `pkg_resources` module |
| |
| This enables the pkg_resources module to find metadata from wheels |
| that have been included in this .par file. |
| |
| The functions and classes here are scoped to this function, since |
| we might have multitple pkg_resources modules, or none. |
| """ |
| |
| try: |
| __import__(pkg_resources_name) |
| pkg_resources = sys.modules.get(pkg_resources_name) |
| if pkg_resources is None: |
| return |
| except ImportError: |
| # Skip setup |
| return |
| |
| if not _version_check_pkg_resources(pkg_resources): |
| # Skip setup |
| return |
| |
| class DistInfoMetadata(pkg_resources.EggMetadata): |
| """Metadata provider for zip files containing .dist-info |
| |
| In find_dist_info_in_zip(), we call |
| metadata.resource_listdir(directory_name). However, it doesn't |
| work with EggMetadata, because _zipinfo_name() expects the |
| directory name to end with a /, but metadata._listdir() which |
| expects the directory to _not_ end with a /. |
| |
| Therefore this class exists. |
| """ |
| |
| def _zipinfo_name(self, fspath): |
| """Overrides EggMetadata._zipinfo_name""" |
| # Convert a virtual filename (full path to file) into a |
| # zipfile subpath usable with the zipimport directory |
| # cache for our target archive |
| fspath = fspath.rstrip(os.sep) |
| if fspath == self.loader.archive: |
| return '' |
| if fspath.startswith(self.zip_pre): |
| return fspath[len(self.zip_pre):] |
| raise AssertionError( |
| "%s is not a subpath of %s" % (fspath, self.zip_pre) |
| ) |
| |
| def _parts(self, zip_path): |
| """Overrides EggMetadata._parts""" |
| # Convert a zipfile subpath into an egg-relative path part |
| # list. |
| fspath = self.zip_pre + zip_path |
| if fspath == self.egg_root: |
| return [] |
| if fspath.startswith(self.egg_root + os.sep): |
| return fspath[len(self.egg_root) + 1:].split(os.sep) |
| raise AssertionError( |
| "%s is not a subpath of %s" % (fspath, self.egg_root) |
| ) |
| |
| def find_dist_info_in_zip(importer, path_item, only=False): |
| """Find dist-info style metadata in zip files. |
| |
| importer: PEP 302-style Importer object |
| path_item (str): filename or pseudo-filename like: |
| /usr/somedirs/main.par |
| or |
| /usr/somedirs/main.par/pypi__portpicker_1_2_0 |
| only (bool): We ignore the `only` flag because it's not clear |
| what it should actually do in this case. |
| |
| Yields pkg_resources.Distribution objects |
| """ |
| metadata = DistInfoMetadata(importer) |
| for subitem in metadata.resource_listdir('/'): |
| basename, ext = os.path.splitext(subitem) |
| if ext.lower() == '.dist-info': |
| # Parse distribution name |
| match = pkg_resources.EGG_NAME(basename) |
| project_name = 'unknown' |
| if match: |
| project_name = match.group('name') |
| # Create metadata object |
| subpath = os.path.join(path_item, subitem) |
| submeta = DistInfoMetadata( |
| zipimport.zipimporter(path_item)) |
| # Override pkg_resources defaults to avoid |
| # "resource_filename() only supported for .egg, not |
| # .zip" message |
| submeta.egg_name = project_name |
| submeta.egg_info = subpath |
| submeta.egg_root = path_item |
| dist = pkg_resources.Distribution.from_location( |
| path_item, subitem, submeta) |
| yield dist |
| |
| def find_eggs_and_dist_info_in_zip(importer, path_item, only=False): |
| """Chain together our finder and the standard pkg_resources finder |
| |
| For simplicity, and since pkg_resources doesn't provide a public |
| interface to do so, we hardcode the chaining (find_eggs_in_zip). |
| """ |
| # Our finder |
| for dist in find_dist_info_in_zip(importer, path_item, only): |
| yield dist |
| # The standard pkg_resources finder |
| for dist in pkg_resources.find_eggs_in_zip(importer, path_item, only): |
| yield dist |
| return |
| |
| # This overwrites the existing registered finder. |
| pkg_resources.register_finder(zipimport.zipimporter, |
| find_eggs_and_dist_info_in_zip) |
| |
| # Note that the default WorkingSet has already been created, and |
| # there is no public interface to easily refresh/reload it that |
| # doesn't also have a "Don't use this" warning. So we manually |
| # add just the entries we know about to the existing WorkingSet. |
| for entry in sys.path: |
| importer = pkgutil.get_importer(entry) |
| if isinstance(importer, zipimport.zipimporter): |
| for dist in find_dist_info_in_zip(importer, entry, only=True): |
| if isinstance(dist._provider, DistInfoMetadata): |
| pkg_resources.working_set.add(dist, entry, insert=False, |
| replace=True) |
| |
| |
| def _initialize_import_path(import_roots, import_prefix): |
| """Add extra entries to PYTHONPATH so that modules can be imported.""" |
| # We try to match to order of Bazel's stub |
| full_roots = [ |
| os.path.join(import_prefix, import_root) |
| for import_root in import_roots] |
| sys.path[1:1] = full_roots |
| _log('# adding %s to sys.path' % full_roots) |
| |
| |
| def setup(import_roots, zip_safe): |
| """Initialize subpar run-time support |
| |
| Args: |
| import_root (list): subdirs inside .par file to add to the |
| module import path at runtime. |
| zip_safe (bool): If False, extract the .par file contents to a |
| temporary directory, and import everything from |
| that directory. |
| |
| Returns: |
| True if setup was successful, else False |
| """ |
| archive_path = _find_archive() |
| if not archive_path: |
| warnings.warn('Failed to initialize .par file runtime support', |
| UserWarning) |
| return False |
| if sys.path[0] != archive_path: |
| warnings.warn('Failed to initialize .par file runtime support. ' + |
| 'archive_path was %r, sys.path was %r' % ( |
| archive_path, sys.path), |
| UserWarning) |
| return False |
| |
| # Extract files to disk if necessary |
| if not zip_safe: |
| extract_dir = _extract_files(archive_path) |
| # sys.path[0] is the name of the executing .par file. Point |
| # it to the extract directory instead, so that Python searches |
| # there for imports. |
| sys.path[0] = extract_dir |
| import_prefix = extract_dir |
| else: # Import directly from .par file |
| extract_dir = None |
| import_prefix = archive_path |
| |
| # Initialize import path |
| _initialize_import_path(import_roots, import_prefix) |
| |
| # Add hook for package metadata |
| _setup_pkg_resources('pkg_resources') |
| _setup_pkg_resources('pip._vendor.pkg_resources') |
| |
| return True |