re-write sort function

lua/ansible/init.lua

@@ -5,67 +5,117 @@
 -- telescope window to search through all available module names, enter inserts, other key opens doc in split
 
 
-local ts_utils = require("nvim-treesitter.ts_utils")
+local ts = require('nvim-treesitter.parsers')
 
 local M = {}
 
 M.example = function()
-    local lines = {}
+    local bufnr = vim.api.nvim_get_current_buf()
+    local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+    local parser = ts.get_parser(bufnr, 'yaml')
+    local tree = parser:parse()[1]
+    local root = tree:root()
 
-    -- TOOD
+    local objects = {}
-    -- correctly handle comments
-    -- removing all empty lines is back for raw blocks
 
-    -- remove all whitespace, we will create new blocks separated by empty lines after
-    local last_search = vim.fn.getreg('/')
-    vim.api.nvim_command('g/^$/d')
-    vim.fn.setreg('/', last_search)
     local current_line_nr = 1
-    local buf_handle = vim.api.nvim_win_get_buf(0)
+    while current_line_nr <= #lines do
-    local last_line_nr = vim.api.nvim_buf_line_count(buf_handle)
+        local line = lines[current_line_nr]
-    while (current_line_nr <= last_line_nr) do
-        vim.api.nvim_win_set_cursor(0, {current_line_nr, 1})
-        local node = ts_utils.get_node_at_cursor(0)
-        local parent = node:parent()
-        while (parent ~= nil and node:type() ~= 'block_mapping_pair') do
-            node = parent
-            parent = node:parent()
-        end
-        local start_row, start_column, end_row, end_column = node:range()
-        -- print('start_row=', start_row)
-        -- print('end_row=', end_row)
-        -- print('line:', vim.inspect(vim.fn.getline(start_row + 1, end_row + 1)))
 
-        lines[current_line_nr] = vim.fn.getline(start_row + 1, end_row + 1)
+        -- print('current_line_nr', current_line_nr)
+        -- print('line', line)
 
-        current_line_nr = current_line_nr + 1 + (end_row - start_row)
+        if not line or line:match("^%s*$") then
+            -- print('empty line, continuing')
+            current_line_nr = current_line_nr + 1
+            goto continue
         end
 
-    -- sort does not work if the index is not a continuous integer, so lets make sure it is. if we are sure that we dont need the original line number, we can do this from the start in `lines`
+        -- if line:match("^%s*#") then
-    a = {}
+        --     -- print('comment, continuing')
-    for i,v in pairs(lines) do table.insert(a, v) end
+        --     -- todo - how? they are currently part of the nodes, but the previous one - if anything i want them to stick to the next one
+        --     current_line_nr = current_line_nr + 1
+        --     goto continue
+        -- end
+
+        local node = root:named_descendant_for_range(current_line_nr-1, 0, current_line_nr-1, #line)
+        if node and (node:type() == 'block_mapping_pair') then
+            -- local node_lines = vim.treesitter.get_node_text(node, bufnr) -- not using, returns string, we want a table
+            local start_row, _, end_row, _ = node:range()
+            local node_lines = vim.api.nvim_buf_get_lines(bufnr, start_row, end_row + 1, false)
+            -- print('line', line)
+            -- print('node type', node:type())
+            -- print('node range', node:range())
+            -- print('node lines', node_lines)
+
+            table.insert(objects, {
+                sort_line = line,
+                lines = node_lines,
+            })
+            -- set current line to the next line after the node's end, and accouting for 0-indexing
+            current_line_nr = end_row + 2
+            goto continue
+        end
+
+        -- do not get stuck on lines that do not match above
+        print('did not match the line, node type: ', node:type(), ', line: ', line)
+        -- insert it anyway, we do not want to loose lines
+        table.insert(objects, {
+            sort_line = line,
+            lines = {line},
+        })
+        current_line_nr = current_line_nr + 1
+        ::continue::
+    end
+
+    -- print(vim.inspect(objects))
     table.sort(
-        a,
+        objects,
         function(k1, k2)
-            return k1[1] < k2[1]
+            return k1.sort_line < k2.sort_line
         end
     )
-    -- print(vim.inspect(a))
+    -- print(vim.inspect(objects))
 
-    insert_row = 1
+    -- clear the buffer
-    last_prefix = vim.split(a[1][1], '__')[1]
+    vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {})
-    for _, lines in ipairs(a) do
+
-        current_prefix = vim.split(lines[1], '__')[1]
+    -- flatten sorted objects into lines and insert back into the buffer
-        if current_prefix ~= last_prefix then
+    local sorted_lines = {}
-            vim.fn.setline(insert_row, '')
+    local last_prefix = nil
-            insert_row = insert_row + 1
+    for _, obj in ipairs(objects) do
+
+        -- if the prefix switches, insert an empty line before the lines
+        local split_line = vim.split(obj.sort_line, '__')
+        -- if we only get one element back this means the line does not contain '__'
+        -- lets try splitting on '_'
+        -- else the fallback is the whole line
+        if #split_line <= 1 then
+            split_line = vim.split(obj.sort_line, '_')
         end
+        local current_prefix = split_line[1]
+        -- print('current_prefix', current_prefix)
+        -- print('last_prefix', last_prefix)
+
+        if last_prefix == nil then
+            -- handle the first line
             last_prefix = current_prefix
-        for i, line in ipairs(lines) do
+
-            vim.fn.setline(insert_row, line)
+        elseif current_prefix ~= last_prefix then
-            insert_row = insert_row + 1
+            -- print('inserting newline')
+            table.insert(sorted_lines, "")
+            last_prefix = current_prefix
+        end
+
+        -- print('obj', vim.inspect(obj))
+
+        for _, line in ipairs(obj.lines) do
+            table.insert(sorted_lines, line)
         end
     end
+
+    -- print('sorted_lines', vim.inspect(sorted_lines))
+    vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, sorted_lines)
 end