From a3f904b2dac10a0764b4a5acf286e65cb38c2bde Mon Sep 17 00:00:00 2001
From: David Ibia <david.ibia@boxmarshall.com>
Date: Fri, 21 Jun 2024 12:51:44 +0100
Subject: [PATCH] feat(none-ls.lua): add functions CheckPoetryVirtualEnv and
 ResolvePythonEnvironment to handle virtual environments more efficiently

---
 lua/absolute/after/none-ls.lua | 36 +++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/lua/absolute/after/none-ls.lua b/lua/absolute/after/none-ls.lua
index 7ca9f1d..d5e2077 100644
--- a/lua/absolute/after/none-ls.lua
+++ b/lua/absolute/after/none-ls.lua
@@ -1,5 +1,35 @@
 local null_ls = require("null-ls")
 
+
+function CheckPoetryVirtualEnv()
+    local poetry = vim.fn.system("poetry env info -p 2>/dev/null")
+    if poetry == "" then
+        return nil
+    end
+    return poetry
+end
+
+function ResolvePythonEnvironment()
+    -- Get system global python env
+    local python = vim.fn.system("which python3")
+    -- Check of virtual environments exists
+    local virtual = os.getenv("VIRTUAL_ENV") or os.getenv("CONDA_PREFIX") or CheckPoetryVirtualEnv()
+
+    if virtual == nil then
+        -- Check if poetry environment exists
+
+        print("No virtual environment detected, using system python")
+        return {
+            "--python-executable", python
+        }
+    end
+
+    print("Using virtual environment: " .. virtual)
+    return {
+        "--python-executable", virtual .. "/bin/python3"
+    }
+end
+
 null_ls.setup({
     sources = {
         null_ls.builtins.formatting.stylua,
@@ -7,11 +37,7 @@ null_ls.setup({
         null_ls.builtins.formatting.shfmt,
         null_ls.builtins.formatting.black,
         null_ls.builtins.diagnostics.mypy.with({
-            extra_args = function()
-                -- Let's resolve the virtual environment   or conda environment
-                local virtual = os.getenv("VIRTUAL_ENV") or os.getenv("CONDA_PREFIX") or "/usr"
-                return { "--python-executable", virtual .. "/bin/python3" }
-            end,
+            extra_args = ResolvePythonEnvironment()
         }),
     },
 })