tests: unit tests using luassert (#6)

.gitignore

@@ -39,3 +39,5 @@ luac.out
 *.x86_64
 *.hex
 
+# nvim help tags
+tags

doc/pymple.txt

@@ -77,7 +77,12 @@ config = {
     -- the log level to use
     -- (one of "trace", "debug", "info", "warn", "error", "fatal")
     level = "debug",
-  }
+  },
+  -- python options
+  python = {
+  -- the names of virtual environment folders to look out for
+  virtual_env_names = { ".venv" },
+  },
 }
 ```
 

lua/pymple/config.lua

@@ -54,7 +54,12 @@
 ---     -- the log level to use
 ---     -- (one of "trace", "debug", "info", "warn", "error", "fatal")
 ---     level = "debug",
----   }
+---   },
+---   -- python options
+---   python = {
+---   -- the names of virtual environment folders to look out for
+---   virtual_env_names = { ".venv" },
+---   },
 --- }
 --- ```
 --- </pre>
@@ -106,14 +111,31 @@ local default_logging_options = {
   level = "debug",
 }
 
----@alias Config { keymaps: Keymaps, create_user_commands: UserCommandOptions, update_imports: UpdateImportsOptions, logging: LoggingOptions}
+---@alias PythonOptions { virtual_env_names: string[] }
 
----@type Config
-M.default_config = {
+---@type PythonOptions
+local default_python_options = {
+  virtual_env_names = { ".venv" },
+}
+
+---@alias Config { keymaps: Keymaps, create_user_commands: UserCommandOptions, update_imports: UpdateImportsOptions, logging: LoggingOptions, python: PythonOptions}
+
+local default_config = {
   keymaps = default_keymaps,
   create_user_commands = default_user_command_options,
   update_imports = default_update_imports_options,
   logging = default_logging_options,
+  python = default_python_options,
 }
 
+---@type Config
+M.default_config = default_config
+
+---@param opts Config
+function M.set_config(opts)
+  M.config = vim.tbl_deep_extend("force", default_config, opts)
+end
+
+M.set_config(M.default_config)
+
 return M

lua/pymple/init.lua

@@ -26,6 +26,7 @@ local function setup(opts)
   opts = opts or {}
 
   setmetatable(opts, { __index = config.default_config })
+  config.set_config(opts)
 
   if opts.logging.enabled then
     opts.logging.enabled = nil

lua/pymple/resolve_imports.lua

@@ -19,15 +19,19 @@ local IMPORTABLE_SYMBOLS_PATTERNS = {
 
 ---@param args table
 ---@param symbol string
+---@param regexes string[]
 ---@return table
-local function add_symbol_regexes(args, symbol)
-  for _, pattern in ipairs(IMPORTABLE_SYMBOLS_PATTERNS) do
+local function add_symbol_regexes(args, symbol, regexes)
+  regexes = regexes or IMPORTABLE_SYMBOLS_PATTERNS
+  for _, pattern in ipairs(regexes) do
     table.insert(args, "-e")
     table.insert(args, string.format(pattern, symbol))
   end
   return args
 end
 
+M.add_symbol_regexes = add_symbol_regexes
+
 ---@param symbol string: the symbol for which to resolve an import
 ---@return string[] | nil: list of candidates
 function M.resolve_python_import(symbol)

lua/pymple/update_imports.lua

@@ -20,6 +20,8 @@ local function build_filetypes_args(filetypes)
   return args
 end
 
+M.build_filetypes_args = build_filetypes_args
+
 --[[
 from path.to.dir import file
 from path.to.dir import (
@@ -31,7 +33,8 @@ from path.to.file import Something
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for
-local function update_imports_split(source, destination, filetypes)
+local function update_imports_split(source, destination, filetypes, _job)
+  local __job = _job or jobs.gg_into_sed
   local cwd = vim.fn.getcwd()
 
   local source_relative = Path:new(source):make_relative(cwd)
@@ -71,10 +74,12 @@ local function update_imports_split(source, destination, filetypes)
       .. utils.escape_import_path(destination_module_name)
       .. "/'",
   }
-  jobs.gg_into_sed(gg_args_split, sed_args_base, true)
-  jobs.gg_into_sed(gg_args_split, sed_args_module, true)
+  __job(gg_args_split, sed_args_base, true)
+  __job(gg_args_split, sed_args_module, true)
 end
 
+M.update_imports_split = update_imports_split
+
 --[[
 from path.to.dir import file
 from path.to.dir import (
@@ -85,7 +90,8 @@ from path.to.dir import (
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for
-local function update_imports_monolithic(source, destination, filetypes)
+local function update_imports_monolithic(source, destination, filetypes, _job)
+  local __job = _job or jobs.gg_into_sed
   local cwd = vim.fn.getcwd()
 
   -- path/to/here
@@ -102,21 +108,16 @@ local function update_imports_monolithic(source, destination, filetypes)
     string.format("'%s[\\.\\s]'", utils.escape_import_path(source_import_path)),
     ".",
   }
-  local sed_args = {
-    "'s/"
-      .. utils.escape_import_path(source_import_path)
-      .. "/"
-      .. utils.escape_import_path(destination_import_path)
-      .. "/'",
-  }
   local sed_args = string.format(
     "'s/%s\\([\\. ]\\)/%s\\1/'",
     utils.escape_import_path(source_import_path),
     utils.escape_import_path(destination_import_path)
   )
-  jobs.gg_into_sed(gg_args, { sed_args }, false)
+  __job(gg_args, { sed_args }, false)
 end
 
+M.update_imports_monolithic = update_imports_monolithic
+
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for

lua/pymple/utils.lua

@@ -1,11 +1,11 @@
 M = {}
 
 local filetype = require("plenary.filetype")
-local config = require("pymple.config")
+local cfg = require("pymple.config")
 local log = require("pymple.log")
 
 ---@type number: The time to wait before refreshing open buffers
-DEFAULT_HANG_TIME = 1000
+local DEFAULT_HANG_TIME = 1000
 -- @param hang_time number: The time to wait before refreshing the buffers
 function M.async_refresh_buffers(hang_time)
   vim.defer_fn(function()
@@ -28,25 +28,29 @@ function M.check_binary_installed(binary_name)
   return 1 == vim.fn.executable(binary_name)
 end
 
+---Converts a path to an import 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$", "")
   return result
 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
+---@return string | nil, string: The base path and the last part of the import path
 function M.split_import_on_last_separator(import_path)
   local base_path, module_name = import_path:match("(.-)%.?([^%.]+)$")
   return base_path, module_name
 end
 
+---Escapes a string to be used in a regex
 ---@param import_path string: The import path to be escaped
 function M.escape_import_path(import_path)
   return import_path:gsub("%.", [[\.]])
 end
 
+---Checks if a file is a python file
 ---@param path string: The path to the file
 ---@return boolean: Whether or not the file is a python file
 local function is_python_file(path)
@@ -55,6 +59,7 @@ end
 
 M.is_python_file = is_python_file
 
+---Recursively checks if a directory contains python files
 ---@param path string: The path to the directory
 ---@return boolean: Whether or not the directory contains python files
 local function recursive_dir_contains_python_files(path)
@@ -74,14 +79,14 @@ end
 
 M.recursive_dir_contains_python_files = recursive_dir_contains_python_files
 
----@param buf number: The buffer number
+---Finds the end line number of a docstring in a list of lines
+---@param lines string[]: The lines to search for a docstring
 ---@return number | nil: The height (in lines) of the docstring
-local function find_docstring_end_line_number(buf)
-  local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+local function find_docstring_end_line_number_in_lines(lines)
   -- if the first line does not contain a docstring, return 0
   if not lines[1]:match('^"""') then
     return 0
-  elseif lines[1]:match('"""$') then
+  elseif lines[1]:match('""".*"""$') then
     return 1
   else
     for i = 2, #lines do
@@ -93,35 +98,66 @@ local function find_docstring_end_line_number(buf)
   return nil
 end
 
+M.find_docstring_end_line_number_in_lines =
+  find_docstring_end_line_number_in_lines
+
+---Finds the end line number of a file docstring
+---@param buf number: The buffer number
+---@return number | nil: The height (in lines) of the docstring
+local function find_docstring_end_line_number(buf)
+  local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+  return find_docstring_end_line_number_in_lines(lines)
+end
+
+M.find_docstring_end_line_number = find_docstring_end_line_number
+
+---Adds an import to the current buffer
 ---@param import_path string: The import path to be added
 ---@param symbol string: The symbol to be imported
-function M.add_import_to_current_buf(import_path, symbol)
-  local docstring_height = find_docstring_end_line_number(0)
-  local insert_on_line
-  if docstring_height == 0 then
-    insert_on_line = 0
-  else
+---@param buf number: The buffer number
+function M.add_import_to_buffer(import_path, symbol, buf)
+  local docstring_height = find_docstring_end_line_number(buf)
+  local insert_on_line = 0
+  if docstring_height ~= 0 then
+    -- add 2 to the docstring height to account for the empty line after the docstring
     insert_on_line = docstring_height + 1
   end
   vim.api.nvim_buf_set_lines(
-    0,
+    buf or 0,
     insert_on_line,
     insert_on_line,
     false,
-    { "from " .. import_path .. " import " .. symbol }
+    { "from " .. import_path .. " import " .. symbol, "" }
   )
 end
 
+---@param path string: The path in which to search for a virtual environment
+---@return string | nil: The path to the virtual environment, or nil if it doesn't exist
+local function dir_contains_virtualenv(path)
+  for _, venv_name in ipairs(cfg.config.python.virtual_env_names) do
+    local venv_path = path .. "/" .. venv_name
+    if vim.fn.isdirectory(venv_path) == 1 then
+      return venv_path
+    end
+  end
+  return nil
+end
+
 ---Get the path to the current virtual environment, or nil if we can't find one
+---@param from_path string: The path to start searching from
 ---@return string | nil: The path to the current virtual environment
-function M.get_virtual_environment()
+function M.get_virtual_environment(from_path)
   local venv = os.getenv("VIRTUAL_ENV")
   if venv then
     return venv
   end
-  local venv_path = vim.fn.getcwd() .. "/.venv"
-  if vim.fn.isdirectory(venv_path) == 1 then
-    return venv_path
+  local current_path = from_path
+  while current_path ~= vim.fn.expand("~") do
+    local venv_path = dir_contains_virtualenv(current_path)
+    if venv_path then
+      return venv_path
+    end
+    current_path = vim.fn.fnamemodify(current_path, ":h")
   end
   return nil
 end
@@ -166,14 +202,14 @@ M.print_msg = print_msg
 ---Print an error message to the console
 ---@param err_msg string: The error message to print
 function M.print_err(err_msg)
-  print_msg(err_msg, config.HL_GROUPS.Error)
+  print_msg(err_msg, cfg.HL_GROUPS.Error)
   log.error(err_msg)
 end
 
 ---Print an info message to the console
 ---@param info_msg string: The info message to print
 function M.print_info(info_msg)
-  print_msg(info_msg, config.HL_GROUPS.More)
+  print_msg(info_msg, cfg.HL_GROUPS.More)
   log.info(info_msg)
 end
 

lua/tests/fixtures/utils/docstrings/misc_1.py

@@ -0,0 +1,4 @@
+"""this is a weirdly formatted docstring
+"""
+
+1 + 1 == 2

lua/tests/fixtures/utils/docstrings/misc_2.py

@@ -0,0 +1,4 @@
+"""
+this is a weirdly formatted docstring"""
+
+1 + 1 == 2

lua/tests/fixtures/utils/docstrings/misc_3.py

@@ -0,0 +1,6 @@
+"""this 
+is a 
+
+weirdly formatted docstring"""
+
+1 + 1 == 2

lua/tests/fixtures/utils/docstrings/multiline_docstring.py

@@ -0,0 +1,9 @@
+"""
+This is a
+multi line docstring
+"""
+
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass

lua/tests/fixtures/utils/docstrings/no_docstring.py

@@ -0,0 +1,4 @@
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass

lua/tests/fixtures/utils/docstrings/single_line_docstring.py

@@ -0,0 +1,6 @@
+"""This is a single line docstring"""
+
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass

lua/tests/resolve_imports_spec.lua

@@ -0,0 +1,16 @@
+local resolve_imports = require("pymple.resolve_imports")
+
+describe("add_symbol_regexes", function()
+  it("std", function()
+    local args = {}
+    local symbol = "foo"
+    local regexes = { "a %sb", "b %sa" }
+    local result = resolve_imports.add_symbol_regexes(args, symbol, regexes)
+    assert.are.same({
+      "-e",
+      "a foob",
+      "-e",
+      "b fooa",
+    }, result)
+  end)
+end)