diff --git a/runtime/support.py b/runtime/support.py index 83991d6..0e89192 100644 --- a/runtime/support.py +++ b/runtime/support.py @@ -33,8 +33,10 @@ """ import os +import pkgutil import sys import warnings +import zipimport def _log(msg): @@ -61,10 +63,21 @@ def _find_archive(): return archive_path -def _setup_pkg_resources(): +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 - import zipimport + __import__(pkg_resources_name) + pkg_resources = sys.modules.get(pkg_resources_name) + if pkg_resources is None: + return except ImportError: # Skip setup return @@ -125,17 +138,20 @@ def find_eggs_and_dist_info_in_zip(importer, path_item, only=False): yield dist return - # This overwrites the existing registered finder. - # - # Note that this also doesn't update the default WorkingSet created by - # pkg_resources when it is imported, since there is no public - # interface to do so that doesn't also have a "Don't use this" - # warning. 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): + pkg_resources.working_set.add(dist, entry, insert=False, + replace=False) def setup(import_roots=None): @@ -154,4 +170,5 @@ def setup(import_roots=None): sys.path.insert(1, new_path) # Add hook for package metadata - _setup_pkg_resources() + _setup_pkg_resources('pkg_resources') + _setup_pkg_resources('pip._vendor.pkg_resources') diff --git a/runtime/support_test.py b/runtime/support_test.py index ab25bd6..2a75334 100644 --- a/runtime/support_test.py +++ b/runtime/support_test.py @@ -21,7 +21,7 @@ class SupportTest(unittest.TestCase): - def test_log(self): + def test__log(self): old_stderr = sys.stderr try: mock_stderr = io.StringIO() @@ -36,17 +36,34 @@ def test_log(self): finally: sys.stderr = old_stderr - def test_find_archive(self): + def test__find_archive(self): # pylint: disable=protected-access path = support._find_archive() self.assertNotEqual(path, None) def test_setup(self): - support.setup(import_roots=['some_root', 'another_root']) - self.assertTrue(sys.path[1].endswith('subpar/runtime/some_root'), - sys.path) - self.assertTrue(sys.path[2].endswith('subpar/runtime/another_root'), - sys.path) + # `import pip` can cause arbitrary sys.path changes, + # especially if using the Debian `python-pip` package or + # similar. Get that lunacy out of the way before starting + # test + try: + import pip + except ImportError: + pass + + old_sys_path = sys.path + try: + mock_sys_path = list(sys.path) + sys.path = mock_sys_path + support.setup(import_roots=['some_root', 'another_root']) + finally: + sys.path = old_sys_path + self.assertTrue(mock_sys_path[1].endswith('subpar/runtime/some_root'), + mock_sys_path) + self.assertTrue(mock_sys_path[2].endswith('subpar/runtime/another_root'), + mock_sys_path) + self.assertEqual(mock_sys_path[0], sys.path[0]) + self.assertEqual(mock_sys_path[3:], sys.path[1:]) if __name__ == '__main__': diff --git a/tests/package_pkg_resources/main.py b/tests/package_pkg_resources/main.py index b37ae54..2888bc2 100644 --- a/tests/package_pkg_resources/main.py +++ b/tests/package_pkg_resources/main.py @@ -26,7 +26,7 @@ def main(): print('Skipping test, pkg_resources module is not available') return - ws = pkg_resources.WorkingSet() + ws = pkg_resources.working_set # Informational for debugging distributions = list(ws)