From a49c7308a5f064f970bfaf7accd7a772452d69fd Mon Sep 17 00:00:00 2001
From: saucoide <32314353+saucoide@users.noreply.github.com>
Date: Wed, 2 Oct 2024 06:15:33 +0200
Subject: [PATCH] feat: support for uv-installed pythons (#842)

* dummy support for uv python installs

* attempt install before confirming the interpreter is available

* adding some tests

* change version output to json so it's consistent in linux vs windows

* skip test with uv-python-support in windows

* fix using the wrong variable

caused problems with pypy

* change min supported uv version

* remove if statement

* change global variable and handle uv missing in uv_version

* typo
---
 nox/virtualenv.py        | 41 ++++++++++++++++++++++++++++++++
 tests/test_virtualenv.py | 51 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/nox/virtualenv.py b/nox/virtualenv.py
index adf60333..aa5f3edf 100644
--- a/nox/virtualenv.py
+++ b/nox/virtualenv.py
@@ -17,6 +17,7 @@
 import abc
 import contextlib
 import functools
+import json
 import os
 import platform
 import re
@@ -28,6 +29,8 @@
 from socket import gethostbyname
 from typing import Any, ClassVar
 
+from packaging import version
+
 import nox
 import nox.command
 from nox.logger import logger
@@ -65,7 +68,39 @@ def find_uv() -> tuple[bool, str]:
     return uv_on_path is not None, "uv"
 
 
+def uv_version() -> version.Version:
+    """Returns uv's version defaulting to 0.0 if uv is not available"""
+    try:
+        ret = subprocess.run(
+            [UV, "version", "--output-format", "json"],
+            check=False,
+            text=True,
+            capture_output=True,
+        )
+    except FileNotFoundError:
+        logger.info("uv binary not found.")
+        return version.Version("0.0")
+
+    if ret.returncode == 0 and ret.stdout:
+        return version.Version(json.loads(ret.stdout).get("version"))
+    else:
+        logger.info("Failed to establish uv's version.")
+        return version.Version("0.0")
+
+
+def uv_install_python(python_version: str) -> bool:
+    """Attempts to install a given python version with uv"""
+    ret = subprocess.run(
+        [UV, "python", "install", python_version],
+        check=False,
+    )
+    return ret.returncode == 0
+
+
 HAS_UV, UV = find_uv()
+# supported since uv 0.3 but 0.4.16 is the first version that doesn't cause
+# issues for nox with pypy/cpython confusion
+UV_PYTHON_SUPPORT = uv_version() >= version.Version("0.4.16")
 
 
 class InterpreterNotFound(OSError):
@@ -526,6 +561,12 @@ def _resolved_interpreter(self) -> str:
             self._resolved = cleaned_interpreter
             return self._resolved
 
+        if HAS_UV and UV_PYTHON_SUPPORT:
+            uv_python_success = uv_install_python(cleaned_interpreter)
+            if uv_python_success:
+                self._resolved = cleaned_interpreter
+                return self._resolved
+
         # The rest of this is only applicable to Windows, so if we don't have
         # an interpreter by now, raise.
         if _SYSTEM != "Windows":
diff --git a/tests/test_virtualenv.py b/tests/test_virtualenv.py
index fd5bfc90..962f3e93 100644
--- a/tests/test_virtualenv.py
+++ b/tests/test_virtualenv.py
@@ -269,6 +269,11 @@ def test_uv_creation(make_one):
     assert venv._check_reused_environment_type()
 
 
+@has_uv
+def test_uv_managed_python(make_one):
+    make_one(interpreter="cpython3.12", venv_backend="uv")
+
+
 def test_constructor_defaults(make_one):
     venv, _ = make_one()
     assert venv.location
@@ -620,6 +625,50 @@ def find_uv_bin():
     assert nox.virtualenv.find_uv() == (found, path)
 
 
+@pytest.mark.parametrize(
+    ["return_code", "stdout", "expected_result"],
+    [
+        (0, '{"version": "0.2.3", "commit_info": null}', "0.2.3"),
+        (1, None, "0.0"),
+        (1, '{"version": "9.9.9", "commit_info": null}', "0.0"),
+    ],
+)
+def test_uv_version(monkeypatch, return_code, stdout, expected_result):
+    def mock_run(*args, **kwargs):
+        return subprocess.CompletedProcess(
+            args=["uv", "version", "--output-format", "json"],
+            stdout=stdout,
+            returncode=return_code,
+        )
+
+    monkeypatch.setattr(subprocess, "run", mock_run)
+    assert nox.virtualenv.uv_version() == version.Version(expected_result)
+
+
+def test_uv_version_no_uv(monkeypatch):
+    def mock_exception(*args, **kwargs):
+        raise FileNotFoundError
+
+    monkeypatch.setattr(subprocess, "run", mock_exception)
+    assert nox.virtualenv.uv_version() == version.Version("0.0")
+
+
+@pytest.mark.parametrize(
+    ["requested_python", "expected_result"],
+    [
+        ("3.11", True),
+        ("pypy3.8", True),
+        ("cpython3.9", True),
+        ("python3.12", True),
+        ("nonpython9.22", False),
+        ("java11", False),
+    ],
+)
+@has_uv
+def test_uv_install(requested_python, expected_result):
+    assert nox.virtualenv.uv_install_python(requested_python) == expected_result
+
+
 def test_create_reuse_venv_environment(make_one, monkeypatch):
     # Making the reuse requirement more strict
     monkeypatch.setenv("NOX_ENABLE_STALENESS_CHECK", "1")
@@ -840,6 +889,7 @@ def special_run(cmd, *args, **kwargs):
 
 
 @mock.patch("nox.virtualenv._SYSTEM", new="Windows")
+@mock.patch("nox.virtualenv.UV_PYTHON_SUPPORT", new=False)
 def test__resolved_interpreter_windows_path_and_version(make_one, patch_sysfind):
     # Establish that if we get a standard pythonX.Y path, we look it
     # up via the path on Windows.
@@ -865,6 +915,7 @@ def test__resolved_interpreter_windows_path_and_version(make_one, patch_sysfind)
 @pytest.mark.parametrize("sysfind_result", [r"c:\python37-x64\python.exe", None])
 @pytest.mark.parametrize("sysexec_result", ["3.7.3\\n", RAISE_ERROR])
 @mock.patch("nox.virtualenv._SYSTEM", new="Windows")
+@mock.patch("nox.virtualenv.UV_PYTHON_SUPPORT", new=False)
 def test__resolved_interpreter_windows_path_and_version_fails(
     input_, sysfind_result, sysexec_result, make_one, patch_sysfind
 ):