feat: smarter project discovery

lua/pymple/utils.lua

@@ -62,14 +62,19 @@ local MAX_UPWARD_JUMPS = 5
 ---@param starting_dir string | nil: The directory to start searching from
 ---@return string | nil: The root of the python project
 local function find_project_root(starting_dir, root_markers)
+  local pythonpath = os.getenv("PYTHONPATH")
+  if pythonpath then
+    return pythonpath
+  end
+
   local dir = starting_dir or vim.fn.getcwd()
   local jumps = 0
-  while dir ~= "/" do
-    if jumps > MAX_UPWARD_JUMPS then
-      return nil
-    end
+  while dir ~= "/" and jumps <= MAX_UPWARD_JUMPS do
     for _, marker in ipairs(root_markers) do
       if vim.fn.glob(dir .. "/" .. marker) ~= "" then
+        if vim.fn.glob(dir .. "/src") ~= "" then
+          dir = dir .. "/src"
+        end
         return dir
       end
     end

lua/tests/utils_spec.lua

@@ -4,9 +4,26 @@ local cwd = vim.fn.getcwd()
 local mock = require("luassert.mock")
 
 describe("find_project_root", function()
-  it("std", function()
+  it("pythonpath", function()
+    mock(os, "getenv", function()
+      return "/some/path"
+    end)
+    local result = utils.find_project_root(nil, {})
+    assert.equals("/some/path", result)
+    mock.revert(os)
+  end)
+
+  it("no pythonpath and src", function()
+    local result = utils.find_project_root(
+      FIXTURES_PATH .. "/project_with_src/src",
+      { "pyproject.toml" }
+    )
+    assert.equals(FIXTURES_PATH .. "/project_with_src/src", result)
+  end)
+
+  it("no pythonpath and no src", function()
     local result = utils.find_project_root(
-      FIXTURES_PATH .. "/project/src",
+      FIXTURES_PATH .. "/project/module/__init__.py",
       { "pyproject.toml" }
     )
     assert.equals(FIXTURES_PATH .. "/project", result)
@@ -217,7 +234,6 @@ describe("get_virtual_environment", function()
       .. "/virtual_environments/present"
     local result = utils.get_virtual_environment(working_dir, { ".venv" })
     assert.equals(working_dir .. "/.venv", result)
-    mock.revert(vim.fn)
   end)
 
   it("no venv", function()