From 2593945f6320b8f70ff34580fd20c60529795ba5 Mon Sep 17 00:00:00 2001
From: Alexandre Feblot <alexandre.feblot@gmail.com>
Date: Sun, 29 Sep 2024 16:31:32 +0200
Subject: [PATCH] Add support to detect uv based projects

---
 autoswitch_virtualenv.plugin.zsh | 18 ++++++++++++++
 tests/test_get_venv_type.zunit   | 27 ++++++++++++++++++++
 tests/test_mkvenv.zunit          | 42 ++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)

diff --git a/autoswitch_virtualenv.plugin.zsh b/autoswitch_virtualenv.plugin.zsh
index cec9d9f..5279fba 100644
--- a/autoswitch_virtualenv.plugin.zsh
+++ b/autoswitch_virtualenv.plugin.zsh
@@ -56,6 +56,10 @@ function _get_venv_type() {
         venv_type="poetry"
     elif [[ -f "$venv_dir/requirements.txt" || -f "$venv_dir/setup.py" ]]; then
         venv_type="virtualenv"
+    elif [[ -f "$venv_dir/uv.lock" ]]; then
+        venv_type="uv"
+    elif [[ -f "$venv_dir/pyproject.toml" && $(grep -c '[tool.uv]' "$venv_dir/pyproject.toml") -gt 0 ]]; then
+        venv_type="uv"
     fi
     printf "%s" "$venv_type"
 }
@@ -160,6 +164,12 @@ function _activate_pipenv() {
 }
 
 
+function _activate_uv() {
+    _maybeworkon ".venv" "virtualenv"
+    return 0
+}
+
+
 # Automatically switch virtualenv when $AUTOSWITCH_FILE file detected
 function check_venv()
 {
@@ -318,6 +328,14 @@ function mkvenv()
         poetry install $params
         _activate_poetry
         return
+    elif [[ "$venv_type" == "uv" ]]; then
+        if ! type "uv" > /dev/null; then
+            _missing_error_message uv
+            return
+        fi
+        uv sync $params
+        _activate_uv
+        return
     else
         if ! type "virtualenv" > /dev/null; then
             _missing_error_message virtualenv
diff --git a/tests/test_get_venv_type.zunit b/tests/test_get_venv_type.zunit
index f8e7b90..ece55cc 100644
--- a/tests/test_get_venv_type.zunit
+++ b/tests/test_get_venv_type.zunit
@@ -27,6 +27,33 @@
     assert "$output" same_as "poetry"
 }
 
+@test '_get_venv_type uv not found' {
+    touch "$TARGET/pyproject.toml"
+
+    run _get_venv_type "$TARGET"
+
+    assert $state equals 0
+    assert "$output" same_as "virtualenv"
+}
+
+@test '_get_venv_type uv found from pyproject.toml' {
+    echo "[uv.tool]" > "$TARGET/pyproject.toml"
+
+    run _get_venv_type "$TARGET"
+
+    assert $state equals 0
+    assert "$output" same_as "uv"
+}
+
+@test '_get_venv_type uv found from uv.lock' {
+    touch "$TARGET/uv.lock"
+
+    run _get_venv_type "$TARGET"
+
+    assert $state equals 0
+    assert "$output" same_as "uv"
+}
+
 @test '_get_venv_type virtualenv (requirements.txt)' {
     touch "$TARGET/requirements.txt"
 
diff --git a/tests/test_mkvenv.zunit b/tests/test_mkvenv.zunit
index 39882d5..642472b 100644
--- a/tests/test_mkvenv.zunit
+++ b/tests/test_mkvenv.zunit
@@ -26,6 +26,14 @@
     function _activate_pipenv {
         echo "activating pipenv"
     }
+
+    function uv {
+        echo "uv" $@
+    }
+
+    function _activate_uv {
+        echo "activating uv"
+    }
 }
 
 @teardown {
@@ -81,6 +89,18 @@
     assert "$lines[2]" same_as "activating pipenv"
 }
 
+@test 'mkvenv - (uv project) runs correct command' {
+    mkdir myproject
+    cd myproject
+    touch "uv.lock"
+
+    run mkvenv
+
+    assert $status equals 0
+    assert "$lines[1]" same_as "uv sync"
+    assert "$lines[2]" same_as "activating uv"
+}
+
 @test 'mkvenv - uses default python if set and not specified' {
     mkdir myproject
     cd myproject
@@ -189,3 +209,25 @@
     assert "$lines[5]" is_empty
     assert ${#lines} equals 4
 }
+
+@test 'prints help message and disables plugin if uv not setup' {
+    touch "uv.lock"
+
+    # Mock type to fail
+    function type() {
+        if [[ "$1" == "uv" ]]; then
+            return 1
+        fi
+        return 0
+    }
+
+    run mkvenv
+
+    assert $status equals 0
+    assert "$lines[1]" contains "zsh-autoswitch-virtualenv requires 'uv' to install this project!"
+    assert "$lines[2]" is_empty
+    assert "$lines[3]" contains "If this is already installed but you are still seeing this message,"
+    assert "$lines[4]" contains "then make sure the \e[1muv\e[0m command is in your PATH.\n"
+    assert "$lines[5]" is_empty
+    assert ${#lines} equals 4
+}