Add modules support into resolve_imports (#47)

Co-authored-by: Alexandre Pasmantier <alex.pasmant@gmail.com>

Author: SavingFrame

lua/pymple/api.lua

@@ -58,20 +58,17 @@ M.resolve_import_under_cursor = function()
     )
     return
   elseif #candidates == 1 then
-    local final_import = candidates[1]
+    local final_import = utils.to_import_statement(candidates[1])
     utils.add_import_to_buffer(
-      "from " .. final_import .. " import " .. symbol,
+      final_import,
       0,
       config.user_config.add_import_to_buf.autosave
     )
-    log.debug("Added import for " .. symbol .. ": " .. final_import)
   else
     local candidate_statements = {}
     for _, candidate in ipairs(candidates) do
-      table.insert(
-        candidate_statements,
-        "from " .. candidate .. " import " .. symbol
-      )
+      local import_statement = utils.to_import_statement(candidate)
+      table.insert(candidate_statements, import_statement)
     end
     local longest_candidate = utils.longest_string_in_list(candidate_statements)
     local telescope_opts = {}

lua/pymple/build.lua

@@ -64,6 +64,32 @@ local function install_sed(os_name)
   end
 end
 
+local function install_fd(os_name)
+  local out, code = nil, nil
+  if os_name == "Darwin" then
+    log.debug("MacOS detected, installing fd...")
+    local job = Job:new({
+      command = "brew",
+      args = { "install", "fd" },
+    })
+    out, code = job:sync(JOB_TIMEOUT)
+  elseif os_name == "Linux" then
+    log.debug("Linux detected, installing sed...")
+    -- make this depend on the package manager
+    local job = Job:new({
+      command = "sudo",
+      args = { "apt-get", "install", "fd-find" },
+    })
+    out, code = job:sync(JOB_TIMEOUT)
+  end
+  if is_outcode_success(code) then
+    return true
+  else
+    log.warning(out)
+    return false
+  end
+end
+
 function M.build()
   -- check if cargo is installed
   if not is_cargo_installed() then
@@ -98,6 +124,16 @@ function M.build()
     utils.print_err(message)
     return
   end
+
+  -- install fd depending on OS
+  utils.print_info("Installing sed/gsed...")
+  if install_fd(utils.OS_NAME) then
+    utils.print_info("fd installed successfully")
+  else
+    local message = "Failed to install fd. Try installing it manually."
+    utils.print_err(message)
+    return
+  end
 end
 
 return M

lua/pymple/health.lua

@@ -26,6 +26,12 @@ local required_binaries = {
     binaries = { "gg" },
     min_version = "0.2.23",
   },
+  {
+    name = "fd",
+    url = "[sharkdp/fd](https://github.com/sharkdp/fd)",
+    optional = false,
+    binaries = { "fd" },
+  },
   {
     name = "sed",
     url = "[https://www.gnu.org/software/sed](https://www.gnu.org/software/sed/manual/sed.html)",

lua/pymple/resolve_imports/init.lua

@@ -88,23 +88,43 @@ function M.resolve_python_import(symbol, current_file_path)
     .. current_file_path
     .. " "
     .. build_symbol_regexes(symbol, IMPORTABLE_SYMBOLS_PATTERNS)
+  --- fd -H -I -p "/sqlalchemy/__init__.py$|/sqlalchemy.py$" ~.venv/lib/python3.11/site-package
+  local fd_args = '-H -I -p "/'
+    .. symbol
+    .. "/__init__.py$|/"
+    .. symbol
+    .. ".py$"
+    .. '"'
+    .. " "
   local candidate_paths = jobs.find_import_candidates(gg_args, target_paths)
+  local modules_paths =
+    jobs.find_import_modules_candidates(fd_args, target_paths)
+  log.debug("Modules: " .. vim.inspect(modules_paths))
   log.debug("Candidates: " .. vim.inspect(candidate_paths))
   candidate_paths = filter_candidates(candidate_paths)
+  modules_paths = filter_candidates(modules_paths)
+  local result = {}
 
-  local import_candidates = {}
   for _, path in ipairs(candidate_paths) do
     local _path = utils.make_relative_to(path, root)
-    local import_path = utils.to_import_path(_path)
-    table.insert(import_candidates, import_path)
+    local reference_path = utils.to_python_reference_path(_path, symbol)
+    table.insert(result, reference_path)
+  end
+
+  for _, path in ipairs(modules_paths) do
+    local _path = utils.make_relative_to(path, root)
+    local reference_path = utils.to_python_reference_path(_path)
+    table.insert(result, reference_path)
   end
-  -- sort imports by length
-  table.sort(import_candidates, function(a, b)
+
+  result = utils.deduplicate_list(result)
+
+  table.sort(result, function(a, b)
     -- TODO: add other sorting rules (e.g. alphabetical, private modules, etc.)
     return #a < #b
   end)
 
-  return import_candidates
+  return result
 end
 
 return M

lua/pymple/resolve_imports/jobs.lua

@@ -26,4 +26,26 @@ function M.find_import_candidates(gg_args, sys_paths)
   return candidates
 end
 
+function M.find_import_modules_candidates(fd_args, sys_paths)
+  local candidates = {}
+  for _, sys_path in ipairs(sys_paths) do
+    local fd_command = "fd " .. fd_args .. sys_path
+    log.debug("FD command: " .. fd_command)
+    local job = vim.system({ utils.SHELL, "-c", fd_command }):wait()
+    local result_lines = vim.split(job.stdout, "\n")
+    for i, line in ipairs(result_lines) do
+      if line == "" then
+        table.remove(result_lines, i)
+      end
+    end
+    if #result_lines ~= 0 then
+      for _, line in ipairs(result_lines) do
+        local path = utils.make_relative_to(line, sys_path)
+        table.insert(candidates, path)
+      end
+    end
+  end
+  return candidates
+end
+
 return M

lua/pymple/update_imports/init.lua

@@ -55,7 +55,7 @@ M.make_monolithic_imports_job = make_monolithic_imports_job
 function M.prepare_jobs(source, destination, filetypes, python_root)
   local s, d = unpack(
     utils.map(
-      utils.to_import_path,
+      utils.to_python_reference_path,
       utils.make_files_relative({ source, destination }, python_root),
       {}
     )

lua/pymple/utils.lua

@@ -97,14 +97,37 @@ end
 
 M.find_project_root = find_project_root
 
----Converts a path to an import path
+---Converts a path to a python reference path
 ---@param module_path string: The path to a python module
----@return string: The import path for the module
-function M.to_import_path(module_path)
-  local result, _ = module_path:gsub("/", "."):gsub("%.py$", "")
+---@param python_symbol string | nil: The symbol to reference
+---@return string: The python reference path
+function M.to_python_reference_path(module_path, python_symbol)
+  local result, _ =
+    module_path:gsub("/", "."):gsub("%.py$", ""):gsub(".__init__", "")
+  if python_symbol then
+    result = result .. "." .. python_symbol
+  end
   return result
 end
 
+---Generate import string from reference path
+---@param reference_path string: The reference path
+---@return string: The import string
+function M.to_import_statement(reference_path)
+  -- Split, but keep "__init__" as part of the path
+  local parts = vim.split(reference_path, ".", { plain = true })
+  -- If there's only one part, this is a simple import
+  if #parts == 1 then
+    return "import " .. parts[1]
+  end
+  local module_name = parts[#parts]
+  table.remove(parts, #parts)
+
+  local module_path = table.concat(parts, ".")
+
+  return "from " .. module_path .. " import " .. module_name
+end
+
 ---Splits an import path on the last separator
 ---@param import_path string: The import path to be split
 ---@return string, string: The base path and the last part of the import path

lua/tests/utils_spec.lua

@@ -29,11 +29,41 @@ describe("find_project_root", function()
   end)
 end)
 
-describe("to_import_path", function()
+describe("to_python_reference_path", function()
   it("std", function()
-    local result = utils.to_import_path("foo/bar/baz.py")
+    local result = utils.to_python_reference_path("foo/bar/baz.py")
     assert.equals("foo.bar.baz", result)
   end)
+
+  it("__init__.py", function()
+    local result = utils.to_python_reference_path("foo/bar/__init__.py")
+    assert.equals("foo.bar", result)
+  end)
+
+  it("python_symbol", function()
+    local result = utils.to_python_reference_path("foo/bar/baz.py", "qux")
+    assert.equals("foo.bar.baz.qux", result)
+  end)
+
+  it("__init__ and python_symbol", function()
+    local result = utils.to_python_reference_path("foo/bar/__init__.py", "qux")
+    assert.equals("foo.bar.qux", result)
+  end)
+end)
+
+describe("to_import_statement", function()
+  it("std", function()
+    local result = utils.to_import_statement("foo.bar.baz")
+    assert.equals("from foo.bar import baz", result)
+  end)
+  it("python_symbol", function()
+    local result = utils.to_import_statement("foo.bar.baz.qux")
+    assert.equals("from foo.bar.baz import qux", result)
+  end)
+  it("single_symbol", function()
+    local result = utils.to_import_statement("foo")
+    assert.equals("import foo", result)
+  end)
 end)
 
 describe("split_import_on_last_separator", function()