Migrating vim-coiled-snake from Vim to Neovim broke a couple things; and I improved my previous folding preferences from before.
The two things I cared about: syntax highlighting on folded lines, and having decorated classes and functions being folded properly.
Here's the updated script.
-- Custom foldtext: skip @decorator lines, show class/def signature with TS highlights.
vim.opt_local.foldtext = "v:lua._python_foldtext()"
_G._python_foldtext = function()
return _G._python_foldtext_impl(vim.v.foldstart, vim.v.foldend)
end
function _G._python_foldtext_impl(start_lnum, end_lnum)
local bufnr = vim.api.nvim_get_current_buf()
local nlines = end_lnum - start_lnum + 1
-- Skip @decorator lines, find first class/def line
local chosen = start_lnum
for lnum = start_lnum, math.min(start_lnum + 5, end_lnum) do
if not vim.fn.getline(lnum):match("^%s*@") then
chosen = lnum
break
end
end
local line = vim.fn.getline(chosen)
local indent = string.rep(" ", vim.fn.indent(chosen))
-- Per-char Treesitter captures for syntax-highlighted chunks
local chunks = { { indent, "Normal" } }
local prev_hl = "Normal"
local buf = {}
for col = 0, #line - 1 do
local hl = "Normal"
local caps = vim.treesitter.get_captures_at_pos(bufnr, chosen - 1, col)
if caps and #caps > 0 then
hl = "@" .. caps[1].capture
end
if hl ~= prev_hl then
if #buf > 0 then
table.insert(chunks, { table.concat(buf), prev_hl })
buf = {}
end
prev_hl = hl
end
table.insert(buf, line:sub(col + 1, col + 1))
end
if #buf > 0 then
table.insert(chunks, { table.concat(buf), prev_hl })
end
table.insert(chunks, { string.format(" (%d lines)", nlines), "Comment" })
return chunks
end
vim.opt_local.foldtextonly accepts strings in Neovim 0.11, not functions. You have to usefoldtext = "v:lua._python_foldtext()"as a string expression, notfoldtext = _python_foldtext`.
The other thing that bit me was saved vim-views.
I keep viewoptions = "cursor,folds" because I want folds remembered between sessions.
But Neovim restores saved views after the ftplugin runs, which means the ftplugin's fold settings get overwritten by whatever was saved last time.
vim.schedule runs too early for this.
The only thing that worked was vim.defer_fn with a 50ms delay to re-close imports after the view restore fires.
vim.opt_local.foldlevel = 0
vim.defer_fn(function()
for lnum = 1, vim.fn.line("$") do
local line = vim.fn.getline(lnum)
if not (line:match("^%s*import%s") or line:match("^%s*from%s.+%simport%s")) then
if vim.fn.foldlevel(lnum) > 0 and vim.fn.foldclosed(lnum) ~= -1 then
vim.cmd(lnum .. "foldopen")
end
else
if vim.fn.foldlevel(lnum) > 0 and vim.fn.foldclosed(lnum) == -1 then
vim.cmd(lnum .. "foldclose")
end
end
end
end, 50)
foldlevel = 0 collapses everything initially, then the deferred loop opens anything that isn't an import.
The 50ms delay is enough for the view restore to complete first.
It's tricky to iterate upon so here is the test suite:
-- E2E test for Python folding with Coiled Snake
-- Run: nvim --headless "+luafile /tmp/test_python_folding.lua"
local failures = 0
local passed = 0
local function check(cond, msg)
if cond then
passed = passed + 1
print("PASS: " .. msg)
else
failures = failures + 1
print("FAIL: " .. msg)
end
end
local test_file = "/tmp/test_py_folding_e2e.py"
local f = io.open(test_file, "w")
f:write([[
import os
import sys
import json
import subprocess
import pathlib
import typing
import collections
import functools
import itertools
import dataclasses
import abc
import io
import re
import time
import logging
import traceback
import warnings
import inspect
import contextlib
import urllib
import http
import socket
import select
import struct
import ctypes
import threading
import multiprocessing
import queue
import asyncio
import concurrent
@dataclass
class Config:
name: str
value: int
class User:
def __init__(self, id, email):
self.id = id
def main():
pass
main()
]])
f:close()
vim.cmd("edit " .. test_file)
vim.cmd("sleep 500m")
check(vim.opt_local.foldmethod:get() == "expr", "foldmethod=expr")
check(vim.opt_local.foldexpr:get() == "coiledsnake#FoldExpr(v:lnum)", "foldexpr=coiledsnake")
check(vim.opt_local.foldenable:get() == true, "foldenable=true")
local first_import = nil
for i = 1, vim.fn.line("$") do
if vim.fn.getline(i):match("^import ") then first_import = i break end
end
check(first_import ~= nil, "found first import")
if first_import then
check(vim.fn.foldclosed(first_import) ~= -1, "top-level imports folded by default")
end
local all_imports_folded = true
for i = 1, 31 do
local l = vim.fn.getline(i)
if (l:match("^import ") or l:match("^from ")) and vim.fn.foldclosed(i) == -1 then
all_imports_folded = false break
end
end
check(all_imports_folded, "all import lines within closed folds")
local class_folded = false
for i = 1, vim.fn.line("$") do
local l = vim.fn.getline(i)
if (l:match("^class ") or l:match("^def ")) and vim.fn.foldclosed(i) ~= -1 then
class_folded = true break
end
end
check(not class_folded, "classes/functions NOT folded by default")
local class_line = nil
for i = 1, vim.fn.line("$") do
if vim.fn.getline(i):match("^class Config") then class_line = i break end
end
if class_line then
local fold_start = class_line
if vim.fn.getline(class_line - 1):match("^%s*@") then fold_start = class_line - 1 end
vim.cmd(tostring(fold_start))
vim.cmd("normal! zc")
check(vim.fn.foldclosed(fold_start) == fold_start, "can manually fold @dataclass class")
local ft = _G._python_foldtext_impl(fold_start, vim.fn.foldclosedend(fold_start))
local text = ""
for _, c in ipairs(ft) do text = text .. c[1] end
check(text:match("class Config") ~= nil, "foldtext shows 'class Config' not @dataclass")
check(text:match("@dataclass") == nil, "foldtext does NOT show @dataclass")
local has_ts_hl = false
for _, c in ipairs(ft) do
if c[2]:match("^@") then has_ts_hl = true break end
end
check(has_ts_hl, "foldtext preserves TS syntax highlight groups")
vim.cmd("normal! zo")
end
local func_line = nil
for i = 1, vim.fn.line("$") do
if vim.fn.getline(i):match("^def main") then func_line = i break end
end
if func_line then
vim.cmd(tostring(func_line))
vim.cmd("normal! zc")
check(vim.fn.foldclosed(func_line) == func_line, "can manually fold plain function")
vim.cmd("normal! zo")
end
check(vim.opt_local.foldtext:get() == "v:lua._python_foldtext()", "foldtext uses v:lua._python_foldtext()")
os.remove(test_file)
print(string.format("\n=== %d passed, %d failed ===", passed, failures))
if failures > 0 then vim.cmd("cquit 1") else vim.cmd("qa") end