├── README.md ├── autoload └── nvim_treesitter_pyfold.vim ├── images ├── Classes_Folding.png ├── Function_based_on_doc.png ├── Functions_example.png ├── Multiline_signature.png ├── folding.png └── listlikes.png ├── lua ├── nvim-treesitter-pyfold.lua └── nvim-treesitter-pyfold │ └── folding.lua ├── plugin └── nvim-treesitter-pyfold.vim └── queries └── python └── folds.scm /README.md: -------------------------------------------------------------------------------- 1 | # nvim-treesitter-pyfold 2 | Folding is a great way to declutter the screen from irrelevant context but 3 | overfolding can be equally annoying. 4 | Hence, a lua based module for nvim-treesitter, setting some sane folding 5 | defaults and provides an optional foldtext. 6 | 7 | ## Installation 8 | Using your favourite package manager 9 | ```lua 10 | # vim-plug 11 | Plug 'nvim-treesitter-pyfold' 12 | 13 | # Packer 14 | require('packer').setup { 15 | ... 16 | use 'eddiebergman/nvim-treesitter-pyfold' 17 | ... 18 | } 19 | ``` 20 | 21 | Once your package manager has installed the package, in your nvim-treesitter 22 | setup, you simply enable the package with: 23 | ```lua 24 | require('nvim-treesitter.configs').setup { 25 | ... 26 | pyfold = { 27 | enable = true 28 | custom_foldtext = true -- Sets provided foldtext on window where module is active 29 | } 30 | ... 31 | } 32 | ``` 33 | 34 | If you havn't already, to use treesitter folding, you have to set the following two lines in your `.vimrc` at some point. 35 | ``` 36 | set foldmethod=expr 37 | set foldexpr=nvim_treesitter#fold_expr() 38 | ``` 39 | 40 | For debuging issues, make these two lines are set using `set foldmethod` and `set foldexpr` to verify these are set. 41 | 42 | ## Example 43 | #### Classes 44 | ![Classes folding](https://github.com/eddiebergman/nvim-treesitter-pyfold/blob/master/images/Classes_Folding.png) 45 | #### Functions / Methods 46 | ![Functions example](https://github.com/eddiebergman/nvim-treesitter-pyfold/blob/master/images/Functions_example.png) 47 | ![Multline signature](https://github.com/eddiebergman/nvim-treesitter-pyfold/blob/master/images/Multiline_signature.png) 48 | ![Functions Based on Docs present](https://github.com/eddiebergman/nvim-treesitter-pyfold/blob/master/images/Function_based_on_doc.png) 49 | #### Dicts, Lists, Tuples 50 | ![Dicts, Lists, Tuples folding](https://github.com/eddiebergman/nvim-treesitter-pyfold/blob/master/images/listlikes.png) 51 | 52 | ### TODO 53 | * Getting multiple imports to fold into one line. 54 | ```Python 55 | # Code 56 | import one 57 | import two 58 | import three 59 | from x import y 60 | import a as b 61 | from z import ( 62 | alpha, beta, phi 63 | ) 64 | 65 | # Desired 66 | -- imports 67 | ``` 68 | 69 | The C part of the python tree-sitter parser doesn't want to fold any 70 | 'queryable' that only lives on one line, even if you group multiple of them 71 | together in query. Please let me know if you know of any workaround for this! 72 | -------------------------------------------------------------------------------- /autoload/nvim_treesitter_pyfold.vim: -------------------------------------------------------------------------------- 1 | function! nvim_treesitter_pyfold#foldtext() 2 | return luaeval(printf("require('nvim-treesitter-pyfold.folding').foldtext(%d, %d, %d)", v:foldstart, v:foldend, v:folddashes)) 3 | endfunction 4 | -------------------------------------------------------------------------------- /images/Classes_Folding.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/Classes_Folding.png -------------------------------------------------------------------------------- /images/Function_based_on_doc.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/Function_based_on_doc.png -------------------------------------------------------------------------------- /images/Functions_example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/Functions_example.png -------------------------------------------------------------------------------- /images/Multiline_signature.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/Multiline_signature.png -------------------------------------------------------------------------------- /images/folding.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/folding.png -------------------------------------------------------------------------------- /images/listlikes.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eddiebergman/nvim-treesitter-pyfold/f0d31fc70e8f8fed43ef745814317eab653454de/images/listlikes.png -------------------------------------------------------------------------------- /lua/nvim-treesitter-pyfold.lua: -------------------------------------------------------------------------------- 1 | local M = {} 2 | 3 | M.is_supported = function (lang) return lang == 'python' end 4 | 5 | function M.init() 6 | require('nvim-treesitter').define_modules({ 7 | pyfold = { 8 | enable = true, 9 | disable = {}, 10 | custom_foldtext = true, 11 | module_path = "nvim-treesitter-pyfold.folding", 12 | is_supported = M.is_supported 13 | } 14 | }) 15 | end 16 | 17 | 18 | return M 19 | -------------------------------------------------------------------------------- /lua/nvim-treesitter-pyfold/folding.lua: -------------------------------------------------------------------------------- 1 | local queries = require('vim.treesitter.query') 2 | local configs = require('nvim-treesitter.configs') 3 | local invalidate_query_cache = require('nvim-treesitter.query').invalidate_query_cache 4 | 5 | local api = vim.api 6 | local fn = vim.fn 7 | local exec = vim.api.nvim_exec 8 | 9 | local rtp_search_str = '**/nvim-treesitter-pyfold/queries/python/folds.scm' 10 | 11 | local M = { 12 | cache = {} 13 | } 14 | 15 | local function folds_path() 16 | if M.cache['fold_path'] == nil then 17 | M.cache['fold_path'] = api.nvim_get_runtime_file(rtp_search_str, false)[1] 18 | end 19 | 20 | return M.cache['fold_path'] 21 | end 22 | 23 | local function readfile(path) 24 | local f = io.open(path, 'r') 25 | local content = f:read('*a') 26 | f:close() 27 | return content 28 | end 29 | 30 | function M.attach(bufnr, lang) 31 | local config = configs.get_module('pyfold') 32 | 33 | -- Set fold regions 34 | local fold_query = readfile(folds_path()) 35 | 36 | local version = vim.version() 37 | fold_query = fold_query:gsub('fold', 'foldopen') 38 | if version and version.major == 0 and version.minor >= 9 then 39 | queries.set('python', 'folds', fold_query) 40 | else 41 | queries.set_query('python', 'folds', fold_query) 42 | end 43 | 44 | -- Change to custom foldtext 45 | if config.custom_foldtext == true then 46 | M.prev_foldtext = exec('echo &foldtext', true) 47 | vim.wo.foldtext = 'nvim_treesitter_pyfold#foldtext()' 48 | end 49 | 50 | end 51 | 52 | function M.detach(bufnr) 53 | local config = configs.get_module('pyfold') 54 | invalidate_query_cache('python', 'folds') 55 | 56 | if config.custom_foldtext == true then 57 | vim.wo.foldtext = M.prev_foldtext 58 | end 59 | end 60 | 61 | function M.is_supported(lang) 62 | return lang == 'python' 63 | end 64 | 65 | local function is_doc_fold(s, e) 66 | return s:find('"""') ~= nil and e:find('"""') ~= nil 67 | end 68 | 69 | local function is_doc_and_body(s, e) 70 | return s:find('"""') ~= nil and e:find('"""') == nil 71 | end 72 | 73 | local function is_main_func(s, e) 74 | return s:find('__main__') ~= nil and s:find('__name__') ~= nil 75 | end 76 | 77 | local function is_dict(s, e) 78 | return s:find('{%s*$') ~= nil and e:find('}%s*$') ~= nil 79 | end 80 | 81 | local function is_list(s, e) 82 | return s:find('%[%s*$') ~= nil and e:find('%]%s*$') ~= nil 83 | end 84 | 85 | local function is_tuple(s, e) 86 | return s:find('%(%s*$') ~= nil and e:find('%)%s*$') ~= nil 87 | end 88 | 89 | function M.foldtext(lstart, lend, dashes) 90 | local s = fn.getline(lstart) 91 | local e = fn.getline(lend) 92 | 93 | if is_doc_fold(s, e) then 94 | -- replace """ with |, if nothing after | on same line, 95 | -- replace that with | doc 96 | local s2 = s 97 | if lstart ~= lend and s:find('"""%s*$') ~= nil then 98 | s2 = s:gsub('"""', 'o ') .. fn.getline(lstart+1):match("^%s*(.-)%s*$") 99 | end 100 | return s2:gsub('"""%s*', 'o ' ):gsub('o%s*$', 'o doc') 101 | 102 | elseif is_doc_and_body(s, e) then 103 | -- replace """ with |, if noting after | on same line, 104 | -- replace that with "| doc, body" 105 | -- 106 | local s2 = s 107 | if lstart ~= lend and s:find('"""%s*$') ~= nil then 108 | s2 = s:gsub('"""', 'o ') .. fn.getline(lstart+1):match("^%s*(.-)%s*$") 109 | end 110 | return s2:gsub('"""%s*', 'o▶ '):gsub('o▶%s*$', 'o▶ doc, body') 111 | 112 | elseif is_main_func(s, e) then 113 | return s 114 | 115 | elseif is_dict(s, e) then 116 | local nlines = tostring(lend - lstart -1) 117 | return s:gsub('{.*$', '{ ... }')..' ('..nlines..')' 118 | 119 | elseif is_list(s, e) then 120 | local nlines = tostring(lend - lstart -1) 121 | return s:gsub('%[.*$', '[ ... ]')..' ('..nlines..')' 122 | 123 | elseif is_tuple(s, e) then 124 | local nlines = tostring(lend - lstart -1) 125 | return s:gsub('%(.*$', '( ... )')..' ('..nlines..')' 126 | 127 | else 128 | return s:gsub('[^%s].*$', '▶ body ') 129 | end 130 | end 131 | 132 | return M 133 | -------------------------------------------------------------------------------- /plugin/nvim-treesitter-pyfold.vim: -------------------------------------------------------------------------------- 1 | if exists('g:loaded_nvim_treesitter_pyfold') | finish | endif 2 | 3 | let s:save_cpo = &cpo 4 | set cpo&vim 5 | 6 | lua require('nvim-treesitter-pyfold').init() 7 | 8 | let &cpo = s:save_cpo 9 | unlet s:save_cpo 10 | 11 | let g:loaded_nvim_treesitter_pyfold = 1 12 | -------------------------------------------------------------------------------- /queries/python/folds.scm: -------------------------------------------------------------------------------- 1 | ; Fold functions bodies with their optional docstring if present 2 | (function_definition 3 | body: (block) @fold 4 | ) 5 | (function_definition 6 | body: (block 7 | (expression_statement (string))? @fold 8 | ) 9 | ) 10 | 11 | ; Fold class definitions bodies with their optional docstring if present 12 | (class_definition 13 | body: (block) @fold 14 | ) 15 | (class_definition 16 | body: (block 17 | (expression_statement (string))? @fold 18 | ) 19 | ) 20 | 21 | ; Fold imports at top TODO: Doesnt work as intended 22 | ;(module . 23 | ; ( 24 | ; [(import_statement) (import_from_statement)] 25 | ; [(import_statement) (import_from_statement)]+ 26 | ; ) @fold 27 | ;) 28 | 29 | ; Fold main 30 | (module 31 | (if_statement 32 | condition : (comparison_operator 33 | (identifier) @id (#match? @id "^__name__") 34 | [(string) (list)] 35 | ) 36 | ) @fold 37 | ) 38 | 39 | ; Fold dicts 40 | (expression_statement 41 | (assignment 42 | right: (dictionary) 43 | ) @fold 44 | ) 45 | 46 | ; Fold lists 47 | (expression_statement 48 | (assignment 49 | right: (list) 50 | ) @fold 51 | ) 52 | 53 | ; Fold Tuples 54 | (expression_statement 55 | (assignment 56 | right: (tuple) 57 | ) @fold 58 | ) 59 | 60 | ; Fold multiline decorators 61 | (decorator 62 | (call 63 | ) @fold 64 | ) 65 | 66 | 67 | ; Old default folds 68 | ;(while_statement (block) @fold) 69 | ;(for_statement (block) @fold) 70 | ;(if_statement (block) @fold) 71 | ;(with_statement (block) @fold) 72 | ;(try_statement (block) @fold) 73 | 74 | ;[ 75 | ; (import_from_statement) 76 | ; (parameters) 77 | ; (argument_list) 78 | ; 79 | ; (parenthesized_expression) 80 | ; (generator_expression) 81 | ; (list_comprehension) 82 | ; (set_comprehension) 83 | ; (dictionary_comprehension) 84 | ; 85 | ; (tuple) 86 | ; (list) 87 | ; (set) 88 | ; (dictionary) 89 | ; 90 | ; (string) 91 | ;] @fold 92 | --------------------------------------------------------------------------------