Skip to content
2 changes: 1 addition & 1 deletion installer/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,7 @@ def bootstrap():
venv.import_package_bundle(
server_site_packages_path, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES"
)
venv.import_package_bundle(system_site_packages_path)
venv.import_package_bundle(system_site_packages_path, priority=True)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What kind of paths are in front here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue was with server site packages. Since they're dynamically inserted into sys path, they would always be prioritized.

venv.import_package_bundle(kernel_site_package_path)

# Phase 7.2: Symlink notebook to jupyter_server for compatibility with
Expand Down
27 changes: 25 additions & 2 deletions installer/module/virtual_environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,15 +81,29 @@ def start_server(self, command: str, cwd: Optional[str] = None) -> ServerProcess
return server_proc

def import_package_bundle(
self, bundle_site_package_path: str, condition_env: Optional[str] = None
self,
bundle_site_package_path: str,
condition_env: Optional[str] = None,
*,
priority: bool = False,
) -> None:
"""
Import a package bundle to the virtual environment.

:param bundle_site_package_path: Absolute path to the 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'.
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.
:raises ValueError: If both condition_env and priority are specified.
"""
if condition_env and priority:
raise ValueError(
"condition_env and priority are mutually exclusive; "
"specify only one of them"
)

pth_file_path = os.path.join(self._site_packages_path, "deepnote.pth")

with open(pth_file_path, "a+", encoding="utf-8") as pth_file:
Expand All @@ -105,6 +119,15 @@ def import_package_bundle(
"Bundle was conditionally imported to the virtual environment "
f"(loads when {condition_env}=true)."
)
elif priority:
# Insert at front of sys.path for higher priority (overrides other bundles)
pth_content = (
f"import sys; sys.path.insert(0, '{bundle_site_package_path}')"
)
pth_file.write(pth_content + "\n")
logger.info(
"Bundle was imported with priority to the virtual environment."
)
else:
pth_file.write(bundle_site_package_path + "\n")
logger.info(
Expand Down
100 changes: 100 additions & 0 deletions tests/unit/test_virtual_environment.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
"""Tests for VirtualEnvironment class."""

from pathlib import Path
from unittest.mock import patch

import pytest

from installer.module.virtual_environment import VirtualEnvironment


class TestImportPackageBundle:
"""Tests for import_package_bundle method."""

@pytest.fixture
def venv(self, tmp_path):
"""Create a VirtualEnvironment with a temporary path."""
venv_path = tmp_path / "venv"
with patch(
"installer.module.virtual_environment.get_current_python_version",
return_value="3.11",
):
venv = VirtualEnvironment(str(venv_path))
# Create the site-packages directory
Path(venv.site_packages_path).mkdir(parents=True, exist_ok=True)
return venv

def test_import_package_bundle_plain_path(self, venv):
"""Test that plain path is appended to .pth file."""
bundle_path = "/some/bundle/site-packages"

venv.import_package_bundle(bundle_path)

pth_file = Path(venv.site_packages_path) / "deepnote.pth"
content = pth_file.read_text()

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, ...)."""
bundle_path = "/server/libs/site-packages"
condition_env = "DEEPNOTE_INCLUDE_SERVER_PACKAGES"

venv.import_package_bundle(bundle_path, condition_env=condition_env)

pth_file = Path(venv.site_packages_path) / "deepnote.pth"
content = pth_file.read_text()

assert "import os, sys" in content
assert f"sys.path.insert(0, '{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, ...)."""
bundle_path = "/usr/local/lib/python3.11/site-packages"

venv.import_package_bundle(bundle_path, priority=True)

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 "os.environ.get" not in content

def test_import_package_bundle_ordering(self, venv):
"""Test .pth file content matches expected order from __main__.py."""
server_libs = "/tmp/python3.11/server-libs/lib/python3.11/site-packages"
system_site = "/usr/local/lib/python3.11/site-packages"
kernel_libs = "/tmp/python3.11/kernel-libs/lib/python3.11/site-packages"

venv.import_package_bundle(
server_libs, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES"
)
venv.import_package_bundle(system_site, priority=True)
venv.import_package_bundle(kernel_libs)

pth_file = Path(venv.site_packages_path) / "deepnote.pth"
lines = pth_file.read_text().splitlines()

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]
# Line 2: system with priority
assert f"sys.path.insert(0, '{system_site}')" in lines[1]
# Line 3: kernel plain path
assert lines[2].strip() == kernel_libs

def test_import_package_bundle_condition_env_and_priority_mutually_exclusive(
self, venv
):
"""Test that condition_env and priority cannot be used together."""
bundle_path = "/some/bundle/site-packages"

with pytest.raises(ValueError, match="mutually exclusive"):
venv.import_package_bundle(
bundle_path,
condition_env="SOME_ENV_VAR",
priority=True,
)