I found the exact requirement was missing, so I ended up making a plugin on my own.
https://github.com/kiyoon/python-import.nvim
It uses all of the below in order to find the best import statement:
- Pre-defined lookup table
- Search project import statements
- Pyright LSP
The install instruction and usage may change over time (it's an early stage of the plugin) so I'd advise you to check the README for the latest instructions.
The plugin is customisable, for example, you can change the lookup table definitions. That said, this plugin may NOT suit all your requirements. But making such plugin on your own is straightforward.
How to make such plugin on Neovim
1. Identify what you want to import
Let's say you want to import the word under your cursor. You can use a utility function like this in lua:
---Get current word in a buffer
---It is aware of the insert mode (move column by -1 if the mode is insert).
---@param winnr integer?
---@return string
local function get_current_word(winnr)
winnr = winnr or vim.api.nvim_get_current_win()
local bufnr = vim.api.nvim_win_get_buf(winnr)
local line, col, mode
vim.api.nvim_win_call(winnr, function()
line = vim.fn.getline "."
col = vim.fn.col "."
mode = vim.fn.mode "."
end)
if mode == "i" then
-- insert mode has cursor one char to the right
col = col - 1
end
local finish = line:find("[^a-zA-Z0-9_]", col)
-- look forward
while finish == col do
col = col + 1
finish = line:find("[^a-zA-Z0-9_]", col)
end
if finish == nil then
finish = #line + 1
end
local start = vim.fn.match(line:sub(1, col), [[\k*$]])
return line:sub(start + 1, finish - 1)
end
Now, you need to convert the word to the import statement. For example, if your cursor is under np, you want to add a statement like import numpy as np. I used a lookup table that has definitions of common imports, so if I don't like the behaviour I can always force it to return what I like. I also used py-tree-sitter to identify all existing import statements in the project, count them and return the number of import statements. This is the example call of the small cli program that does this job.
$ pip install python-import
$ python-import count /path/my-project logging
00002:from my_module import logging
00001:import logging
2. Identify where to import to.
This is very tricky to define and it depends on your project style. For me personally, I'd put them at the end of existing import statements, or if they don't exist, after all docstrings / comments at the beginning of the document.
I used treesitter to identify the current node types (string, comment, etc.), and justify where it belongs. See the functions below:
---Return line after the first comments and docstring.
---It iterates e.g. 50 first lines and obtains treesitter nodes to check the syntax (string or comment)
---@param bufnr integer?
---@param max_lines integer?
---@return integer?
local function find_line_after_module_docstring(bufnr, max_lines)
bufnr = bufnr or vim.api.nvim_get_current_buf()
max_lines = max_lines or 200
local lines = vim.api.nvim_buf_get_lines(bufnr, 0, max_lines, false)
for i, line in ipairs(lines) do
local node = vim.treesitter.get_node { pos = { i - 1, 0 } }
if
node ~= nil
and node:type() ~= "comment"
and node:type() ~= "string"
and node:type() ~= "string_start"
and node:type() ~= "string_content"
and node:type() ~= "string_end"
then
return i
end
end
return nil
end
---Find the first import statement in a python file.
---@param bufnr integer?
---@param max_lines integer?
local function find_line_first_import(bufnr, max_lines)
bufnr = bufnr or vim.api.nvim_get_current_buf()
max_lines = max_lines or 200
local lines = vim.api.nvim_buf_get_lines(bufnr, 0, max_lines, false)
for i, line in ipairs(lines) do
local node = vim.treesitter.get_node { pos = { i - 1, 0 } }
if node ~= nil and (node:type() == "import_statement" or node:type() == "import_from_statement") then
-- additional check whether the node is top-level.
-- if not, it's probably an import inside a function
if node:parent():type() == "module" then
return i
end
end
end
return nil
end
---Find the last import statement in a python file.
---@param bufnr integer?
---@param max_lines integer?
local function find_line_last_import(bufnr, max_lines)
bufnr = bufnr or vim.api.nvim_get_current_buf()
max_lines = max_lines or 200
local lines = vim.api.nvim_buf_get_lines(bufnr, 0, max_lines, false)
-- iterate backwards
for i = #lines, 1, -1 do
local node = vim.treesitter.get_node { pos = { i - 1, 0 } }
if node ~= nil and (node:type() == "import_statement" or node:type() == "import_from_statement") then
-- additional check whether the node is top-level.
-- if not, it's probably an import inside a function
if node:parent():type() == "module" then
return i
end
end
end
return nil
end
3. Add a line
Now that you know where and what to put, simplely call:
local line_number = 1
local import_statements = { "import numpy as np" }
vim.api.nvim_buf_set_lines(0, line_number - 1, line_number - 1, false, import_statements)
to put the text. You don't have to worry about the duplicates and sorting because isort or ruff will fix those!
4. Make it into one function and add a keymap.
This is just an example of a simple function utilising the above helper functions.
local import_as = { np = "numpy" }
local function add_import_statement()
local current_word = get_current_word()
local line_number = find_line_after_module_docstring() -- always at the top module
if import_as[current_word] ~= nil then
local import_statements = { "import " .. import_as[current_word] .. " as " .. current_word }
vim.api.nvim_buf_set_lines(0, line_number - 1, line_number - 1, false, import_statements)
end
end
vim.keymap.set({"i", "n"}, "<M-CR>", add_import_statement)
