Automatically folding Python docstrings in Neovim

[AllanChain]

View Post on Blog

Python docstrings are extremely suitable to be folded, and folding docstrings makes Python code cleaner. This can be achieved in Neovim without plugin, just with built-in foldexpr and tree-sitter.

---

Good Python docstrings are extremely helpful when reading code. However, if I already know what the function does and what all these parameters mean, I don't need to read the docstrings again. They take too much space on the screen, and create a mental burden for me because the functions become long. I realized that Python docstrings are very friendly to folding because the first line is always the summary of the function or class. Therefore, I decided to fold them by default.

It turns out that it is able to achieve this without any plugins, just with foldexpr and tree-sitter. With foldmethod set to expr, the expression from foldexpr is evaluated for each line to determine its fold level. If foldexpr is set to 'v:lua.vim.treesitter.foldexpr()', the fold level will be determined by tree-sitter settings. The default tree-sitter queries for Python can be found here. To automatically fold Python docstrings only, we need to overwrite the default tree-sitter queries and modify foldexpr to make it only work for Python.

Writing simple tree-sitter queries

Tree-sitter is not only in charge of syntax highlighting in Neovim, but also defines the folding behavior. The query syntax is easy to learn. Take the following example

[
  (import_statement)
  (import_from_statement)
]+ @fold

The square brackets [] denotes alternations, similar to the regular expression [abc], and + also has a very similar meaning as in regular expressions. The above query means fold all import and from ... import statements.

But how to write new queries? Neovim has a handy command called InspectTree which will show the parsed syntax tree of the current buffer. For example, the following Python code:

class Test:
    "Test class"

produces:

(module ; [0, 0] - [3, 0]
  (class_definition ; [0, 0] - [1, 16]
    name: (identifier) ; [0, 6] - [0, 10]
    body: (block ; [1, 4] - [1, 16]
      (expression_statement ; [1, 4] - [1, 16]
        (string ; [1, 4] - [1, 16]
          (string_start) ; [1, 4] - [1, 5]
          (string_content) ; [1, 5] - [1, 15]
          (string_end)))))) ; [1, 15] - [1, 16]

The docstring is under class_definition → field body → block → expression_statement → string. Thus, to match docstrings for classes, we use the following query (source):

(class_definition
  body:
    (block
      .
      (expression_statement
        (string) @fold)))

. is the anchor operator, therefore we only match the first expression with string. If you want to explore more from this query, you can use the EditQuery command, paste the above query in, and move your cursor onto the text @fold. If there are matches, they will be shown in the editor.

Similarly, folding the docstring for the functions:

(function_definition
  body:
    (block
      .
      (expression_statement
        (string) @fold)))

And the queries should be placed at ~/.config/nvim/queries/python/folds.scm and will overwrite the default queries.

Config Neovim

Now we can just set

vim.opt.foldexpr = "v:lua.vim.treesitter.foldexpr()"
vim.opt.foldmethod = "expr"

to fold the docstrings. If your code is not folding, maybe the fold is created but expanded, try to press zc to close it.

To make the folding only work in Python, I created a new function in lua/fold.lua to be used in foldexpr, which just checks the file type of the current buffer.

local M = {}

M.foldexpr = function(lnum)
  local ft = vim.api.nvim_get_option_value("filetype", { buf = 0 })
  if ft == "python" then
    return vim.treesitter.foldexpr(lnum)
  else
    return "0"
  end
end

return M

And I set

vim.opt.foldexpr = "v:lua.require'fold'.foldexpr()"

Also, I prefer transparent folding, which keep the syntax highlighting while folding, just like other IDE:

vim.opt.foldtext = ""

Final result