Line data Source code
1 : -- lua/yoda/python_venv.lua
2 : -- Async Python virtual environment detection
3 :
4 17 : local M = {}
5 :
6 17 : local notify = require("yoda-adapters.notification")
7 17 : local lsp_perf = require("yoda.lsp_performance")
8 :
9 : -- Cache for venv detection results (TTL: 5 minutes)
10 17 : local venv_cache = {}
11 17 : local CACHE_TTL = 300000000000 -- 5 minutes in nanoseconds
12 :
13 : --- Check if cached venv is still valid
14 : --- @param root_dir string Project root directory
15 : --- @return string|nil cached_venv Cached venv path or nil
16 : local function get_cached_venv(root_dir)
17 17 : local cached = venv_cache[root_dir]
18 17 : if not cached then
19 16 : return nil
20 : end
21 :
22 1 : local now = vim.uv.hrtime()
23 1 : if now - cached.timestamp < CACHE_TTL then
24 1 : return cached.venv_path
25 : end
26 :
27 : -- Cache expired
28 : venv_cache[root_dir] = nil
29 : return nil
30 17 : end
31 :
32 : --- Cache venv detection result
33 : --- @param root_dir string Project root directory
34 : --- @param venv_path string|nil Virtual environment path
35 : local function cache_venv(root_dir, venv_path)
36 17 : venv_cache[root_dir] = {
37 17 : venv_path = venv_path,
38 17 : timestamp = vim.uv.hrtime(),
39 17 : }
40 34 : end
41 :
42 : --- Build possible venv paths for a project
43 : --- @param root_dir string Project root directory
44 : --- @return table paths List of possible venv paths
45 : local function build_venv_paths(root_dir)
46 : local function join_path(...)
47 102 : return table.concat({ ... }, "/")
48 17 : end
49 :
50 17 : local cwd = vim.fn.getcwd()
51 17 : return {
52 34 : join_path(root_dir, ".venv", "bin", "python"),
53 34 : join_path(root_dir, "venv", "bin", "python"),
54 34 : join_path(root_dir, "env", "bin", "python"),
55 34 : join_path(cwd, ".venv", "bin", "python"),
56 34 : join_path(cwd, "venv", "bin", "python"),
57 34 : join_path(cwd, "env", "bin", "python"),
58 17 : }
59 17 : end
60 :
61 : --- Walk a list of paths using non-blocking libuv fs_stat.
62 : --- Invokes callback(path) with the first executable found, or callback(nil) if none.
63 : --- vim.schedule is used for the final invocation so callers can safely call Neovim APIs.
64 : --- @param paths table List of paths to check
65 : --- @param idx number Current index (start at 1)
66 : --- @param callback function Called with first executable path or nil
67 : local function find_first_executable_async(paths, idx, callback)
68 119 : if idx > #paths then
69 34 : vim.schedule(function()
70 17 : callback(nil)
71 34 : end)
72 17 : return
73 : end
74 :
75 204 : vim.uv.fs_stat(paths[idx], function(err, stat)
76 102 : if not err and stat and stat.type == "file" and (stat.mode % 512 >= 64) then
77 : vim.schedule(function()
78 : callback(paths[idx])
79 : end)
80 : else
81 102 : find_first_executable_async(paths, idx + 1, callback)
82 : end
83 204 : end)
84 119 : end
85 :
86 : --- Detect virtual environment asynchronously
87 : --- @param root_dir string Project root directory
88 : --- @param callback function Callback function(venv_path or nil)
89 17 : function M.detect_venv_async(root_dir, callback)
90 : -- Check cache first
91 17 : local cached = get_cached_venv(root_dir)
92 17 : if cached then
93 : vim.schedule(function()
94 : callback(cached)
95 : end)
96 : return
97 : end
98 :
99 17 : local venv_start_time = vim.uv.hrtime()
100 17 : local possible_paths = build_venv_paths(root_dir)
101 :
102 34 : find_first_executable_async(possible_paths, 1, function(venv_path)
103 : -- Track performance
104 17 : local found = venv_path ~= nil
105 17 : lsp_perf.track_venv_detection(root_dir, venv_start_time, found)
106 :
107 : -- Cache result
108 17 : cache_venv(root_dir, venv_path)
109 :
110 : -- Notify result
111 17 : if venv_path then
112 : notify.notify(string.format("Python LSP: Using venv at %s", venv_path), "info")
113 : else
114 17 : notify.notify("Python LSP: No venv found, using system Python", "info")
115 : end
116 :
117 : -- Execute callback
118 17 : callback(venv_path)
119 34 : end)
120 34 : end
121 :
122 : --- Apply venv to basedpyright LSP clients
123 : --- @param root_dir string Project root directory
124 : --- @param venv_path string Virtual environment path
125 17 : function M.apply_venv_to_lsp(root_dir, venv_path)
126 2 : local clients = vim.lsp.get_clients({ name = "basedpyright" })
127 5 : for _, client in ipairs(clients) do
128 3 : if client.config and client.config.root_dir == root_dir then
129 2 : local settings = client.config and client.config.settings
130 2 : if settings then
131 2 : if settings.basedpyright and settings.basedpyright.analysis then
132 2 : settings.basedpyright.analysis.pythonPath = venv_path
133 : end
134 2 : if settings.python then
135 2 : settings.python.pythonPath = venv_path
136 : end
137 2 : client.notify("workspace/didChangeConfiguration", { settings = settings })
138 : end
139 : end
140 : end
141 19 : end
142 :
143 : --- Detect and apply venv for a project (main entry point)
144 : --- @param root_dir string Project root directory
145 17 : function M.detect_and_apply(root_dir)
146 2 : if not root_dir then
147 1 : return
148 : end
149 :
150 2 : M.detect_venv_async(root_dir, function(venv_path)
151 1 : if venv_path then
152 : M.apply_venv_to_lsp(root_dir, venv_path)
153 : end
154 2 : end)
155 18 : end
156 :
157 : --- Clear venv cache (useful for testing or forcing re-detection)
158 : --- @param root_dir string|nil Specific root dir to clear, or nil for all
159 17 : function M.clear_cache(root_dir)
160 36 : if root_dir then
161 1 : venv_cache[root_dir] = nil
162 : else
163 35 : venv_cache = {}
164 : end
165 53 : end
166 :
167 : --- Get cache statistics (for debugging)
168 : --- @return table stats Cache statistics
169 17 : function M.get_cache_stats()
170 4 : local count = 0
171 4 : local expired = 0
172 4 : local now = vim.uv.hrtime()
173 :
174 6 : for _, cached in pairs(venv_cache) do
175 2 : count = count + 1
176 2 : if now - cached.timestamp >= CACHE_TTL then
177 : expired = expired + 1
178 : end
179 : end
180 :
181 4 : return {
182 4 : total = count,
183 4 : expired = expired,
184 4 : valid = count - expired,
185 4 : ttl_seconds = CACHE_TTL / 1000000000,
186 4 : }
187 17 : end
188 :
189 : --- Setup user commands for venv management
190 17 : function M.setup_commands()
191 6 : vim.api.nvim_create_user_command("PythonVenvDetect", function()
192 : local root_dir = vim.fs.root(0, { "pyproject.toml", "setup.py", "requirements.txt", ".git" })
193 : if not root_dir then
194 : notify.notify("No Python project root detected", "warn")
195 : return
196 : end
197 :
198 : -- Force re-detection by clearing cache
199 : M.clear_cache(root_dir)
200 : M.detect_and_apply(root_dir)
201 3 : end, { desc = "Manually detect and apply Python virtual environment" })
202 :
203 6 : vim.api.nvim_create_user_command("PythonVenvCache", function()
204 : local stats = M.get_cache_stats()
205 : local lines = {
206 : "=== Python Venv Cache ===",
207 : string.format("Total entries: %d", stats.total),
208 : string.format("Valid entries: %d", stats.valid),
209 : string.format("Expired entries: %d", stats.expired),
210 : string.format("TTL: %d seconds", stats.ttl_seconds),
211 : "",
212 : "Cached venvs:",
213 : }
214 :
215 : for root_dir, cached in pairs(venv_cache) do
216 : local age = (vim.uv.hrtime() - cached.timestamp) / 1000000000
217 : local status = age < stats.ttl_seconds and "✓" or "✗"
218 : table.insert(lines, string.format(" %s %s: %s (age: %.1fs)", status, root_dir, cached.venv_path or "none", age))
219 : end
220 :
221 : notify.notify(table.concat(lines, "\n"), "info", { title = "Venv Cache" })
222 3 : end, { desc = "Show Python venv cache status" })
223 :
224 6 : vim.api.nvim_create_user_command("PythonVenvClear", function()
225 : M.clear_cache()
226 : notify.notify("Python venv cache cleared", "info")
227 3 : end, { desc = "Clear Python venv cache" })
228 20 : end
229 :
230 17 : return M
|
