Repository: bluss/pyproject-local-kernel

Issue:
Enable $PATH from .venv for use-venv configurations

Discussion:

Patch:
diff --git a/src/pyproject_local_kernel/__init__.py b/src/pyproject_local_kernel/__init__.py
index b14a67a..d6d4c23 100644
--- a/src/pyproject_local_kernel/__init__.py
+++ b/src/pyproject_local_kernel/__init__.py
@@ -81,7 +81,11 @@ class ProjectDetection:
     config: Config = dataclasses.field(default_factory=Config)
     error_context: str | None = None
 
-    def get_python_cmd(self, allow_fallback=True, allow_hatch_workaround=False) -> t.Sequence[str | Path] | None:
+    def get_python_cmd(self, *args, **kwargs) -> t.Sequence[Path | str] | None:
+        penv = self.resolve(*args, **kwargs)
+        return penv and penv.python_cmd
+
+    def resolve(self, allow_fallback=True, allow_hatch_workaround=False) -> PythonEnvironment | None:
         """
         allow_hatch_workaround: call out to `hatch env find`
         """
@@ -95,23 +99,26 @@ class ProjectDetection:
         # configuration: use-venv
         if use_venv is not None:
             assert self.path is not None
-            return [self.path.parent / get_venv_bin_python(Path(use_venv))]
+            cmd = [self.path.parent / get_venv_bin_python(Path(use_venv))]
+            # virtualenv PATH
+            venv_bin_dir = cmd[0].parent
+            return PythonEnvironment(cmd, venv_bin_dir)
 
         # configuration: python-cmd
         python_cmd = self.config.python_cmd_normalized()
         if python_cmd is not None:
-            return python_cmd
+            return PythonEnvironment(python_cmd)
 
         # project detection
         result = self.kind.python_cmd()
 
         if result is not None:
-            return result
+            return PythonEnvironment(result)
 
         if (allow_fallback and
             self.kind not in (ProjectKind.NoProject, ProjectKind.InvalidData)):
             if fallback := self._fallback_project_kind().python_cmd():
-                return fallback
+                return PythonEnvironment(fallback)
         return None
 
 
@@ -124,6 +131,22 @@ class ProjectDetection:
         return ProjectKind.Unknown
 
 
+@dataclasses.dataclass
+class PythonEnvironment:
+    "A project's python environment"
+    python_cmd: t.Sequence[str | Path]
+    venv_bin_dir: Path | None = None
+
+    def update_environment(self, env: dict[str, t.Any]):
+        "Update environment variables in dict env"
+        if self.venv_bin_dir is None:
+            return
+        path_env = env.get("PATH", os.defpath)
+        path_entries = path_env.split(os.pathsep)
+        if not path_entries or path_entries[0] != self.venv_bin_dir:
+            env["PATH"] = os.pathsep.join([str(self.venv_bin_dir), *path_entries])
+
+
 def find_pyproject_file_from(curdir, basename="pyproject.toml"):
     cwd = Path(curdir).resolve()
     candidate_dirs = [cwd, *cwd.parents]
diff --git a/src/pyproject_local_kernel/main.py b/src/pyproject_local_kernel/main.py
index a85dafc..3805f8e 100644
--- a/src/pyproject_local_kernel/main.py
+++ b/src/pyproject_local_kernel/main.py
@@ -45,7 +45,7 @@ def _setup_logging():
 
 
 async def async_kernel_start(prov: KernelProvisionerBase, args: argparse.Namespace, extra_args: list[str]):
-    await prov.pre_launch()
+    kernel_kws = await prov.pre_launch()
 
     def expand(arg: str):
         "expand variables in argument"
@@ -53,9 +53,10 @@ async def async_kernel_start(prov: KernelProvisionerBase, args: argparse.Namespa
             arg = arg.format(connection_file=args.connection_file)
         return arg
 
-    cmd = [expand(arg) for arg in prov.kernel_spec.argv] + extra_args
+    cmd = kernel_kws.pop("cmd", None)
+    cmd = [expand(arg) for arg in cmd] + extra_args
 
-    kernel_connect_info = await prov.launch_kernel(cmd)
+    kernel_connect_info = await prov.launch_kernel(cmd, **kernel_kws)
     _logger.debug("info=%r", kernel_connect_info)
 
 
diff --git a/src/pyproject_local_kernel/provisioner.py b/src/pyproject_local_kernel/provisioner.py
index fd269ea..467a9c5 100644
--- a/src/pyproject_local_kernel/provisioner.py
+++ b/src/pyproject_local_kernel/provisioner.py
@@ -72,29 +72,34 @@ class PyprojectKernelProvisioner(LocalProvisioner):
         if find_project.kind == ProjectKind.InvalidData:
             raise RuntimeError("\n".join([_MESSAGE_NO_PYPROJECT, f"Reason: {find_project.error_context}"]))
 
-        python_cmd = find_project.get_python_cmd(allow_fallback=True, allow_hatch_workaround=True)
-
-        if python_cmd is None:
+        python_environment = find_project.resolve(allow_hatch_workaround=True)
+        if python_environment is None:
             raise RuntimeError(_MESSAGE_NO_PYPROJECT)
 
+        if python_environment.venv_bin_dir:
+            kwargs["env"] = t.cast(dict, kwargs.get("env", os.environ.copy()))
+            python_environment.update_environment(kwargs["env"])
+
         # convert path to string and update kernel spec argv
-        python_cmd = list(map(str, python_cmd))
+        python_cmd = list(map(str, python_environment.python_cmd))
         kernel_spec.argv[:] = python_cmd + self.python_kernel_args
 
         if self.sanity_check:
             self._python_environment_sanity_check(find_project, python_cmd, cwd, kwargs)
+        return kwargs
 
     async def pre_launch(self, **kwargs) -> t.Dict[str, t.Any]:
         # note: we could raise an exception here and JupyterLab will show the message
         try:
-            self._pplk_pre_launch(**kwargs)
+            new_kwargs = self._pplk_pre_launch(**kwargs)
         except (OSError, RuntimeError) as exc:
             # an error was encountered, run the fallback kernel instead to present the error
             self.kernel_spec.argv[:] = ["pyproject_local_kernel", f"--fallback-kernel={exc}"] + self.python_kernel_args
+            new_kwargs = kwargs
         except Exception:
             raise  # show to user
         self._log_debug("launching kernel from process pid=%d", os.getpid())
-        return await super().pre_launch(**kwargs)
+        return await super().pre_launch(**new_kwargs)
 
     def _python_environment_sanity_check(self, project: ProjectDetection, python_cmd: list[str], cwd: Path, kwargs: dict):
         # skip sanity for uv because it will install ipykernel
@@ -116,6 +121,8 @@ class PyprojectKernelProvisioner(LocalProvisioner):
 
     async def launch_kernel(self, cmd: t.List[str], **kwargs: t.Any) -> KernelConnectionInfo:
         self._log_info("Launching %r", cmd)
+        self._log_debug("PATH=%r", kwargs.get("env", {}).get("PATH"))
+
         try:
             return await super().launch_kernel(cmd, **kwargs)
         except OSError as exc:


Test patch:
diff --git a/tests/test_identify.py b/tests/test_identify.py
index 952863e..441def4 100644
--- a/tests/test_identify.py
+++ b/tests/test_identify.py
@@ -64,14 +64,23 @@ def test_hatch(path, monkeypatch: pytest.MonkeyPatch):
 @pytest.mark.parametrize("path,unix_cmd,win_cmd", [
     ("tests/identify/use-venv", ".myvenv/bin/python", r".myvenv\Scripts\python.exe"),
 ])
-def test_use_venv(path, unix_cmd, win_cmd):
+def test_use_venv(path, unix_cmd: str, win_cmd: str):
     # venv resolves to absolute path
     pd = identify(path)
     base_dir = Path(path).absolute()
     expected = base_dir / (win_cmd if is_windows() else unix_cmd)
-    assert pd.get_python_cmd() == [Path(expected)]
+
+    pres = pd.resolve()
+    assert pres is not None
+    assert pres.python_cmd == [expected]
     assert pd.path is not None and pd.path.name == "pyproject.toml"
 
+    mock_path = os.pathsep.join(["use", "venv"])
+    environment = {"PATH": mock_path}
+    pres.update_environment(environment)
+    path_components = environment["PATH"].split(os.pathsep)
+    assert path_components == [str(expected.parent), *mock_path.split(os.pathsep)]
+
 
 def is_windows():
     return os.name == "nt"