feat(api): Reset the set venv

README.md

@@ -54,21 +54,50 @@ require('swenv.api').auto_venv()
 Using a file named **.venv** in your projects root folder, it will automatically set the
 virtual-env for such environment.
 
+### `post_reset_venv` Function
+
+The `post_reset_venv` function is a callback that can be defined in the `settings`
+table. It is executed after the virtual environment has been reset. This allows for
+additional custom actions to be performed after the environment reset, such as reloading
+configurations, updating UI elements, or triggering other dependent processes.
+
+#### Usage Example
+
+You can define a `post_reset_venv` function in your settings to perform custom actions
+after resetting the virtual environment:
+
+```lua
+return {
+    "AckslD/swenv.nvim",
+    config = function()
+        require("swenv").setup(
+        {
+            post_set_venv = function()
+                vim.cmd [[:LspRestart]]
+            end,
+        }
+    )
+)
+```
+
 #### project.nvim
 
-With [project_nvim](https://github.com/ahmedkhalf/project.nvim) installed swenv.nvim will activate in-project venvs if present.
+With [project_nvim](https://github.com/ahmedkhalf/project.nvim) installed swenv.nvim
+will activate in-project venvs if present.
 
 This integration is skipped if `auto_create_venv` is enabled
 
 #### Auto Create venv
 
 You can have swenv.nvim attempt to create a venv directory and set to it.
 
-swenv.nvim will search up for a file containing dependencies and create a new venv directory set to `auto_create_venv_dir`
+swenv.nvim will search up for a file containing dependencies and create a new venv
+directory set to `auto_create_venv_dir`
 
 Supported install types in order:
 
-- `pdm sync` with `pdm.lock` (does not modify pdm settings. You need to set the venv directory in pdm)
+- `pdm sync` with `pdm.lock` (does not modify pdm settings. You need to set the venv
+  directory in pdm)
 
 > These require the `venv` module in python:
 
@@ -77,14 +106,12 @@ Supported install types in order:
 - `pip install` with `pyproject.toml`
 
 ```lua
-
 require('swenv').setup({
     -- attempt to auto create and set a venv from dependencies
     auto_create_venv = true,
     -- name of venv directory to create if using pip
     auto_create_venv_dir = ".venv"
 })
-
 ```
 
 #### Auto Command
@@ -106,6 +133,14 @@ vim.api.nvim_create_autocmd("FileType", {
 })
 ```
 
+#### Reset Environment
+
+To reset the virtual environment and restore the original Python version, you can call:
+
+```lua
+require('swenv.api').reset_venv()
+```
+
 ## Configuration
 
 Pass a dictionary into `require("swenv").setup()` with callback functions. These are the

lua/swenv/api.lua

@@ -12,7 +12,9 @@ local create = require('swenv.create')
 
 local ORIGINAL_PATH = vim.fn.getenv('PATH')
 local IS_WINDOWS = vim.uv.os_uname().sysname == 'Windows_NT'
-local current_venv = nil
+local ORIGINAL_VIRTUAL_ENV = vim.fn.getenv('VIRTUAL_ENV')
+local ORIGINAL_CONDA_PREFIX = vim.fn.getenv('CONDA_PREFIX')
+local ORIGINAL_CONDA_DEFAULT_ENV = vim.fn.getenv('CONDA_DEFAULT_ENV')
 
 local update_path = function(path)
   local sep
@@ -34,10 +36,10 @@ M.set_venv_path = function(venv)
     vim.fn.setenv('CONDA_PROMPT_MODIFIER', '(' .. venv.name .. ')')
   else
     vim.fn.setenv('VIRTUAL_ENV', venv.path)
+    vim.fn.setenv('VIRTUAL_ENV_PROMPT', '(' .. venv.name .. ')')
   end
 
   current_venv = venv
-  -- TODO: remove old path
   update_path(venv.path)
 
   if settings.post_set_venv then
@@ -98,13 +100,24 @@ local get_venvs_for = function(base_path, source, opts)
   if base_path == nil then
     return venvs
   end
-  local paths = scan_dir(base_path, vim.tbl_extend('force', { depth = 1, only_dirs = true, silent = true }, opts or {}))
-  for _, path in ipairs(paths) do
-    table.insert(venvs, {
-      name = Path:new(path):make_relative(base_path),
-      path = path,
-      source = source,
-    })
+  local success, paths = pcall(
+    scan_dir,
+    tostring(base_path), -- Ensure path is a string
+    vim.tbl_extend('force', { depth = 1, only_dirs = true, silent = true }, opts or {})
+  )
+
+  if not success then
+    return venvs
+  end
+
+  if paths and #paths > 0 then
+    for _, path in ipairs(paths) do
+      table.insert(venvs, {
+        name = Path:new(path):make_relative(tostring(base_path)),
+        path = path,
+        source = source,
+      })
+    end
   end
   return venvs
 end
@@ -160,15 +173,46 @@ local get_pyenv_base_path = function()
   end
 end
 
+local get_project_venv_path = function(project_dir)
+  local venv_path = Path:new(project_dir) / '.venv'
+  if venv_path:exists() then
+    return venv_path
+  end
+  return nil
+end
+
+-- Function to get project virtual environment
+local get_project_venv = function()
+  local venvs = {}
+  local current_dir = vim.fn.getcwd()
+  local project_venv_path = get_project_venv_path(current_dir)
+
+  if project_venv_path then
+    -- Get the project name from the directory
+    local project_name = vim.fn.fnamemodify(current_dir, ':t')
+    table.insert(venvs, {
+      name = project_name,
+      path = tostring(project_venv_path),
+      source = 'venv',
+    })
+  end
+
+  return venvs
+end
+
 M.get_venvs = function(venvs_path)
   local venvs = {}
   vim.list_extend(venvs, get_venvs_for(venvs_path, 'venv'))
   vim.list_extend(venvs, get_venvs_for(get_pixi_base_path(), 'pixi'))
   vim.list_extend(venvs, get_venvs_for(get_conda_base_path(), 'conda'))
   vim.list_extend(venvs, get_conda_base_env())
+  vim.list_extend(venvs, get_project_venv()) -- Add project venv
   vim.list_extend(venvs, get_venvs_for(get_micromamba_base_path(), 'micromamba'))
   vim.list_extend(venvs, get_venvs_for(get_pyenv_base_path(), 'pyenv'))
-  vim.list_extend(venvs, get_venvs_for(get_pyenv_base_path(), 'pyenv', { only_dirs = false }))
+  vim.list_extend(
+    venvs,
+    get_venvs_for(get_pyenv_base_path(), 'pyenv', { only_dirs = false })
+  )
   return venvs
 end
 
@@ -186,6 +230,38 @@ M.pick_venv = function()
   end)
 end
 
+M.reset_venv = function()
+  -- Reset PATH to original value
+  vim.fn.setenv('PATH', ORIGINAL_PATH)
+
+  -- Reset virtual environment variables
+  if ORIGINAL_VIRTUAL_ENV ~= vim.NIL then
+    vim.fn.setenv('VIRTUAL_ENV', ORIGINAL_VIRTUAL_ENV)
+  else
+    vim.fn.setenv('VIRTUAL_ENV', nil)
+  end
+
+  -- Reset Conda environment variables
+  if ORIGINAL_CONDA_PREFIX ~= vim.NIL then
+    vim.fn.setenv('CONDA_PREFIX', ORIGINAL_CONDA_PREFIX)
+  else
+    vim.fn.setenv('CONDA_PREFIX', nil)
+  end
+
+  if ORIGINAL_CONDA_DEFAULT_ENV ~= vim.NIL then
+    vim.fn.setenv('CONDA_DEFAULT_ENV', ORIGINAL_CONDA_DEFAULT_ENV)
+  else
+    vim.fn.setenv('CONDA_DEFAULT_ENV', nil)
+  end
+
+  -- Reset current_venv
+  current_venv = nil
+
+  if settings.post_reset_venv then
+    settings.post_reset_venv()
+  end
+end
+
 M.set_venv = function(name)
   local venvs = settings.get_venvs(settings.venvs_path)
   local closest_match = best_match(venvs, name)
@@ -215,6 +291,17 @@ local auto_venv_project_nvim = function(project_nvim, venvs)
         return
       end
     end
+    -- Check for .venv directory in the project
+    local project_venv_path = get_project_venv_path(project_dir)
+    if project_venv_path then
+      local venv = {
+        name = 'project_venv',
+        path = project_venv_path,
+        source = 'venv',
+      }
+      M.set_venv_path(venv)
+      return
+    end
   end
 end
 

stylua.toml

@@ -0,0 +1,6 @@
+syntax = "Lua51"
+quote_style = "AutoPreferSingle"
+indent_width = 2
+indent_type = "Spaces"
+line_endings = "Unix"
+column_width = 88