diff --git a/installer/__main__.py b/installer/__main__.py index 08c4b0d2..6c72b2c7 100644 --- a/installer/__main__.py +++ b/installer/__main__.py @@ -270,6 +270,7 @@ def bootstrap(): server_site_packages_path = get_server_site_package_path(toolkit_bundle_path) kernel_site_package_path = get_kernel_site_package_path(toolkit_bundle_path) + venv.clear_package_bundles() venv.import_package_bundle( server_site_packages_path, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES" ) diff --git a/installer/module/virtual_environment.py b/installer/module/virtual_environment.py index edf74ae6..c80e156a 100644 --- a/installer/module/virtual_environment.py +++ b/installer/module/virtual_environment.py @@ -86,6 +86,20 @@ def start_server( return server_proc + @property + def _pth_file_path(self) -> str: + """Path to the .pth file that imports package bundles into the venv.""" + return os.path.join(self._site_packages_path, "deepnote.pth") + + def clear_package_bundles(self) -> None: + """ + Remove all previously imported package bundles. + + A reused virtual environment keeps its .pth file, so call this before + importing bundles to avoid stale or duplicated entries. + """ + Path(self._pth_file_path).unlink(missing_ok=True) + def import_package_bundle( self, bundle_site_package_path: str, @@ -100,8 +114,10 @@ def import_package_bundle( :param condition_env: Optional environment variable name. If provided, the bundle will only be loaded when this env var is set to 'true'. Cannot be combined with priority. - :param priority: If True, insert at front of sys.path to override other bundles. - Cannot be combined with condition_env. + :param priority: If True, insert right after the venv's site-packages to + override other bundles, while packages installed into + the venv keep precedence. Cannot be combined with + condition_env. :raises ValueError: If both condition_env and priority are specified. """ if condition_env and priority: @@ -110,14 +126,23 @@ def import_package_bundle( "specify only one of them" ) - pth_file_path = os.path.join(self._site_packages_path, "deepnote.pth") + # Insert right after the venv's own site-packages, so packages the user + # installs into the venv (e.g. an upgraded pip) are not shadowed. + # Paths are compared resolved, since the venv may be reached via symlinks, + # and written with ascii() as site.py may decode .pth files as ASCII. + venv_site_packages = os.path.realpath(self._site_packages_path) + insert_index = ( + "next((i + 1 for i, p in enumerate(sys.path) " + f"if os.path.realpath(p) == {ascii(venv_site_packages)}), 0)" + ) + bundle_literal = ascii(bundle_site_package_path) - with open(pth_file_path, "a+", encoding="utf-8") as pth_file: + with open(self._pth_file_path, "a+", encoding="utf-8") as pth_file: if condition_env: # Write conditional import that checks environment variable pth_content = ( f"import os, sys; " - f"sys.path.insert(0, '{bundle_site_package_path}') " + f"sys.path.insert({insert_index}, {bundle_literal}) " f"if os.environ.get('{condition_env}', '').lower() == 'true' else None" ) pth_file.write(pth_content + "\n") @@ -126,9 +151,10 @@ def import_package_bundle( f"(loads when {condition_env}=true)." ) elif priority: - # Insert at front of sys.path for higher priority (overrides other bundles) + # Higher priority than the other bundles, lower than the venv itself pth_content = ( - f"import sys; sys.path.insert(0, '{bundle_site_package_path}')" + f"import os, sys; " + f"sys.path.insert({insert_index}, {bundle_literal})" ) pth_file.write(pth_content + "\n") logger.info( diff --git a/tests/unit/test_virtual_environment.py b/tests/unit/test_virtual_environment.py index f646495e..10b0f16d 100644 --- a/tests/unit/test_virtual_environment.py +++ b/tests/unit/test_virtual_environment.py @@ -1,5 +1,7 @@ """Tests for VirtualEnvironment class.""" +import os +import sys from pathlib import Path from unittest.mock import patch @@ -36,7 +38,7 @@ def test_import_package_bundle_plain_path(self, venv): assert content == f"{bundle_path}\n" def test_import_package_bundle_with_condition_env(self, venv): - """Test that conditional import uses sys.path.insert(0, ...).""" + """Test that conditional import is inserted after the venv's site-packages.""" bundle_path = "/server/libs/site-packages" condition_env = "DEEPNOTE_INCLUDE_SERVER_PACKAGES" @@ -46,11 +48,12 @@ def test_import_package_bundle_with_condition_env(self, venv): content = pth_file.read_text() assert "import os, sys" in content - assert f"sys.path.insert(0, '{bundle_path}')" in content + assert f"== {ascii(os.path.realpath(venv.site_packages_path))}" in content + assert f"'{bundle_path}')" in content assert f"os.environ.get('{condition_env}', '').lower() == 'true'" in content def test_import_package_bundle_with_priority(self, venv): - """Test that priority=True uses sys.path.insert(0, ...).""" + """Test that priority=True is inserted after the venv's site-packages.""" bundle_path = "/usr/local/lib/python3.11/site-packages" venv.import_package_bundle(bundle_path, priority=True) @@ -58,8 +61,9 @@ def test_import_package_bundle_with_priority(self, venv): pth_file = Path(venv.site_packages_path) / "deepnote.pth" content = pth_file.read_text() - assert "import sys" in content - assert f"sys.path.insert(0, '{bundle_path}')" in content + assert "import os, sys" in content + assert f"== {ascii(os.path.realpath(venv.site_packages_path))}" in content + assert f"'{bundle_path}')" in content assert "os.environ.get" not in content def test_import_package_bundle_ordering(self, venv): @@ -80,12 +84,91 @@ def test_import_package_bundle_ordering(self, venv): assert len(lines) == 3 # Line 1: server-libs conditional assert "DEEPNOTE_INCLUDE_SERVER_PACKAGES" in lines[0] - assert f"sys.path.insert(0, '{server_libs}')" in lines[0] + assert f"'{server_libs}')" in lines[0] # Line 2: system with priority - assert f"sys.path.insert(0, '{system_site}')" in lines[1] + assert f"'{system_site}')" in lines[1] # Line 3: kernel plain path assert lines[2].strip() == kernel_libs + @pytest.mark.parametrize("include_server_packages", ["true", "false"]) + def test_import_package_bundle_resulting_sys_path_order( + self, venv, monkeypatch, include_server_packages + ): + """Test venv site-packages take precedence over system site-packages.""" + server_libs = "/server-libs/site-packages" + system_site = "/usr/local/lib/python3.11/site-packages" + kernel_libs = "/kernel-libs/site-packages" + venv.import_package_bundle( + server_libs, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES" + ) + venv.import_package_bundle(system_site, priority=True) + monkeypatch.setenv("DEEPNOTE_INCLUDE_SERVER_PACKAGES", include_server_packages) + stdlib = "/usr/local/lib/python3.11" + monkeypatch.setattr("sys.path", [stdlib, venv.site_packages_path]) + + lines = (Path(venv.site_packages_path) / "deepnote.pth").read_text() + for line in lines.splitlines(): + exec(line) # noqa: S102 - mimics how site.py runs .pth import lines + sys.path.append(kernel_libs) + + expected = [stdlib, venv.site_packages_path, system_site] + if include_server_packages == "true": + expected.append(server_libs) + assert sys.path == [*expected, kernel_libs] + + def test_clear_package_bundles_removes_stale_entries(self, venv): + """Test that re-importing bundles in a reused venv drops old entries.""" + pth_file = Path(venv.site_packages_path) / "deepnote.pth" + pth_file.write_text("import sys; sys.path.insert(0, '/stale')\n") + + venv.clear_package_bundles() + venv.import_package_bundle("/kernel-libs/site-packages") + + assert pth_file.read_text() == "/kernel-libs/site-packages\n" + + def test_clear_package_bundles_without_pth_file(self, venv): + """Test that clearing works when no bundles were imported yet.""" + venv.clear_package_bundles() + + assert not (Path(venv.site_packages_path) / "deepnote.pth").exists() + + def test_import_package_bundle_escapes_paths(self, venv, monkeypatch): + """Test that paths with quotes and backslashes are kept verbatim.""" + bundle_path = "/odd'dir\\name/site-packages" + venv.import_package_bundle(bundle_path, priority=True) + monkeypatch.setattr("sys.path", [venv.site_packages_path]) + + line = (Path(venv.site_packages_path) / "deepnote.pth").read_text() + exec(line) # noqa: S102 - mimics how site.py runs .pth import lines + + assert sys.path == [venv.site_packages_path, bundle_path] + + def test_import_package_bundle_writes_ascii_only(self, venv): + """Test that non-ASCII paths are escaped so .pth decodes as ASCII.""" + bundle_path = "/bündle/site-packages" + + venv.import_package_bundle(bundle_path, priority=True) + + raw = (Path(venv.site_packages_path) / "deepnote.pth").read_bytes() + assert raw.isascii() + + def test_import_package_bundle_matches_venv_through_symlink( + self, venv, tmp_path, monkeypatch + ): + """Test the bundle lands after the venv when sys.path uses a symlink.""" + bundle_path = "/usr/local/lib/python3.11/site-packages" + venv.import_package_bundle(bundle_path, priority=True) + link = tmp_path / "venv-link" + link.symlink_to(Path(venv.path), target_is_directory=True) + linked_site = str(link / "lib" / "python3.11" / "site-packages") + stdlib = "/usr/local/lib/python3.11" + monkeypatch.setattr("sys.path", [stdlib, linked_site]) + + line = (Path(venv.site_packages_path) / "deepnote.pth").read_text() + exec(line) # noqa: S102 - mimics how site.py runs .pth import lines + + assert sys.path == [stdlib, linked_site, bundle_path] + def test_import_package_bundle_condition_env_and_priority_mutually_exclusive( self, venv ):