Adds ability to toggle system site package inclusion in runtime

docs/index.rst

@@ -89,6 +89,13 @@ SHIV_INTERPRETER
 This is a boolean that bypasses and console_script or entry point baked into your pyz. Useful for
 dropping into an interactive session in the environment of a built cli utility.
 
+SHIV_SYSTEM_SITE_PACKAGES
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This is a boolean that toggles the inclusion of system site packages in the runtime environment
+for your pyz. The default is `False` (don't include system site packages), unless enabled during
+shiv's build process.
+
 SHIV_ENTRY_POINT
 ^^^^^^^^^^^^^^^^
 

src/shiv/bootstrap/__init__.py

@@ -83,9 +83,36 @@ def extract_site_packages(archive, target_path, compile_pyc, compile_workers=0):
     shutil.move(str(target_path_tmp), str(target_path))
 
 
+def patch_sys_path(new_path, system_site_packages):
+    """Insert a new path into sys.path and optionally remove system site-packages
+
+    :param str new_path: The new path to add
+    :param bool system_site_packages: Toggles the inclusion of site packages
+    """
+    if not system_site_packages:
+        sys.path = [
+            p
+            for p in sys.path
+            if Path(p).stem not in ["site-packages", "dist-packages"]
+        ]
+
+    # get sys.path's length
+    length = len(sys.path)
+
+    # Find the first instance of an existing site-packages on sys.path
+    index = _first_sitedir_index() or length
+
+    # append site-packages using the stdlib blessed way of extending path
+    # so as to handle .pth files correctly
+    site.addsitedir(new_path)
+
+    # reorder to place our site-packages before any others found
+    return sys.path[:index] + sys.path[length:] + sys.path[index:length]
+
+
 def _first_sitedir_index():
     for index, part in enumerate(sys.path):
-        if Path(part).stem == "site-packages":
+        if Path(part).stem in ["site-packages", "dist-packages"]:
             return index
 
 
@@ -103,20 +130,11 @@ def bootstrap():
 
     # determine if first run or forcing extract
     if not site_packages.exists() or env.force_extract:
-        extract_site_packages(archive, site_packages.parent, env.compile_pyc, env.compile_workers)
-
-    # get sys.path's length
-    length = len(sys.path)
-
-    # Find the first instance of an existing site-packages on sys.path
-    index = _first_sitedir_index() or length
-
-    # append site-packages using the stdlib blessed way of extending path
-    # so as to handle .pth files correctly
-    site.addsitedir(site_packages)
+        extract_site_packages(
+            archive, site_packages.parent, env.compile_pyc, env.compile_workers
+        )
 
-    # reorder to place our site-packages before any others found
-    sys.path = sys.path[:index] + sys.path[length:] + sys.path[index:length]
+    sys.path = patch_sys_path(site_packages, env.system_site_packages)
 
     # do entry point import and call
     if env.entry_point is not None and env.interpreter is None:

src/shiv/bootstrap/environment.py

@@ -21,6 +21,7 @@ class Environment:
     ROOT = "SHIV_ROOT"
     FORCE_EXTRACT = "SHIV_FORCE_EXTRACT"
     COMPILE_PYC = "SHIV_COMPILE_PYC"
+    SYSTEM_SITE_PACKAGES = "SHIV_SYSTEM_SITE_PACKAGES"
     COMPILE_WORKERS = "SHIV_COMPILE_WORKERS"
 
     def __init__(
@@ -29,13 +30,15 @@ def __init__(
         entry_point=None,
         always_write_cache=False,
         compile_pyc=True,
+        system_site_packages=False,
     ):
         self.build_id = build_id
         self.always_write_cache = always_write_cache
 
         # properties
         self._entry_point = entry_point
         self._compile_pyc = compile_pyc
+        self._system_site_packages = system_site_packages
 
     @classmethod
     def from_json(cls, json_data):
@@ -70,6 +73,12 @@ def force_extract(self):
     def compile_pyc(self):
         return str_bool(os.environ.get(self.COMPILE_PYC, self._compile_pyc))
 
+    @property
+    def system_site_packages(self):
+        return str_bool(
+            os.environ.get(self.SYSTEM_SITE_PACKAGES, self._system_site_packages)
+        )
+
     @property
     def compile_workers(self):
         try:

src/shiv/cli.py

@@ -85,6 +85,11 @@ def copy_bootstrap(bootstrap_target: Path) -> None:
     default=True,
     help="Whether or not to compile pyc files during initial bootstrap.",
 )
+@click.option(
+    "--system-site-packages/--no-system-site-packages",
+    default=False,
+    help="Inherit system site packages during runtime",
+)
 @click.argument("pip_args", nargs=-1, type=click.UNPROCESSED)
 def main(
     output_file: str,
@@ -94,6 +99,7 @@ def main(
     site_packages: Optional[str],
     compressed: bool,
     compile_pyc: bool,
+    system_site_packages: bool,
     pip_args: List[str],
 ) -> None:
     """
@@ -140,7 +146,10 @@ def main(
 
         # create runtime environment metadata
         env = Environment(
-            build_id=str(uuid.uuid4()), entry_point=entry_point, compile_pyc=compile_pyc
+            build_id=str(uuid.uuid4()),
+            entry_point=entry_point,
+            compile_pyc=compile_pyc,
+            system_site_packages=system_site_packages,
         )
 
         Path(working_path, "environment.json").write_text(env.to_json())

test/test_bootstrap.py

@@ -12,7 +12,13 @@
 
 from unittest import mock
 
-from shiv.bootstrap import import_string, current_zipfile, cache_path, _first_sitedir_index
+from shiv.bootstrap import (
+    import_string,
+    current_zipfile,
+    cache_path,
+    _first_sitedir_index,
+    patch_sys_path,
+)
 from shiv.bootstrap.environment import Environment
 
 
@@ -69,6 +75,45 @@ def test_first_sitedir_index(self):
             assert _first_sitedir_index() is None
 
 
+class TestPathPatching:
+    def setup_method(self, method):
+        self.sys_path = [
+            "/usr/lib/python37.zip",
+            "/usr/lib/python3.7",
+            "/usr/lib/python3.7/lib-dynload",
+            "/usr/local/lib/python3.7/dist-packages",
+            "/usr/lib/python3/dist-packages",
+        ]
+        self.site_packages = (
+            "/home/user/.shiv/app_7c32f985-9c69-4fad-bb60-115948078f05/site-packages"
+        )
+
+    def test_patch_path_with_system(self):
+        with mock.patch.object(sys, "path", self.sys_path):
+            with mock.patch("site.addsitedir", new=lambda s: sys.path.append(s)):
+                new_path = patch_sys_path(self.site_packages, True)
+        assert new_path == [
+            "/usr/lib/python37.zip",
+            "/usr/lib/python3.7",
+            "/usr/lib/python3.7/lib-dynload",
+            self.site_packages,
+            "/usr/local/lib/python3.7/dist-packages",
+            "/usr/lib/python3/dist-packages",
+        ]
+
+    def test_patch_path_no_system(self):
+        with mock.patch.object(sys, "path", self.sys_path):
+            with mock.patch("site.addsitedir", new=lambda s: sys.path.append(s)):
+                new_path = patch_sys_path(self.site_packages, False)
+
+        assert new_path == [
+            "/usr/lib/python37.zip",
+            "/usr/lib/python3.7",
+            "/usr/lib/python3.7/lib-dynload",
+            self.site_packages,
+        ]
+
+
 class TestEnvironment:
     def test_overrides(self):
         env = Environment()
@@ -93,6 +138,10 @@ def test_overrides(self):
         with env_var("SHIV_COMPILE_PYC", "False"):
             assert env.compile_pyc is False
 
+        assert env.system_site_packages is False
+        with env_var("SHIV_SYSTEM_SITE_PACKAGES", "True"):
+            assert env.system_site_packages is True
+
         assert env.compile_workers == 0
         with env_var("SHIV_COMPILE_WORKERS", "1"):
             assert env.compile_workers == 1