Add support for list-item cookies

lua/orgmode/org/mappings.lua

@@ -12,6 +12,7 @@ local ts_utils = require('nvim-treesitter.ts_utils')
 local utils = require('orgmode.utils')
 local tree_utils = require('orgmode.utils.treesitter')
 local Headline = require('orgmode.treesitter.headline')
+local List = require('orgmode.treesitter.list')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -162,6 +163,8 @@ function OrgMappings:toggle_checkbox()
   checkbox = checkbox:gsub('%[[%sXx%-]?%]$', new_val)
   local new_line = line:gsub(pattern, checkbox)
   vim.fn.setline('.', new_line)
+  local list = List:new(tree_utils.closest_list())
+  list:update_parent_cookie()
 end
 
 function OrgMappings:timestamp_up_day()

lua/orgmode/treesitter/headline.lua

@@ -120,6 +120,14 @@ function Headline:remove_closed_date()
   tree_utils.set_node_text(dates['CLOSED'], '', true)
 end
 
+function Headline:cookie()
+  local cookie = self:parse('%[%d?/%d?%]')
+  if cookie then
+    return cookie
+  end
+  return self:parse('%[%d?%d?%d?%%%]')
+end
+
 -- @return tsnode, string
 function Headline:parse(pattern)
   local match = ''

lua/orgmode/treesitter/list.lua

@@ -0,0 +1,63 @@
+local ts_utils = require('nvim-treesitter.ts_utils')
+local tree_utils = require('orgmode.utils.treesitter')
+local query = vim.treesitter.query
+local Headline = require('orgmode.treesitter.headline')
+
+local List = {}
+
+function List:new(list_node)
+  local data = { list = list_node }
+  setmetatable(data, self)
+  self.__index = self
+  return data
+end
+
+-- Updates the cookie of the immediate parent
+-- This always checks for a parent list first
+-- then for a headline.
+function List:parent_cookie()
+  local parent_list = tree_utils.find_list(self.list:parent())
+  if parent_list then
+    -- We only care about the cookie if it's at the top
+    local top_item = parent_list:named_child(0)
+    local content = top_item:field('contents')[1]
+    -- The cookie should be the last thing on the line
+    local cookie_node = content:named_child(content:named_child_count() - 1)
+
+    local text = query.get_node_text(cookie_node, 0)
+    if text:match('%[%d?/%d?%]') or text:match('%[%d?%d?%d?%%%]') then
+      return cookie_node
+    end
+  end
+
+  local parent_header = Headline:new(tree_utils.closest_headline())
+  return parent_header:cookie()
+end
+
+function List:update_parent_cookie()
+  local parent_cookie = self:parent_cookie()
+  if not parent_cookie then
+    return
+  end
+
+  local checkboxes = self:checkboxes()
+  local checked_boxes = vim.tbl_filter(function(box)
+    return box:match('%[%w%]')
+  end, checkboxes)
+  local new_status
+  if query.get_node_text(parent_cookie, 0):find('%%') then
+    new_status = ('[%d%%]'):format((#checked_boxes / #checkboxes) * 100)
+  else
+    new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  end
+  tree_utils.set_node_text(parent_cookie, new_status)
+end
+
+function List:checkboxes()
+  return vim.tbl_map(function(node)
+    local text = query.get_node_text(node, 0)
+    return text:match('%[.%]')
+  end, ts_utils.get_named_children(self.list))
+end
+
+return List

lua/orgmode/utils/treesitter.lua

@@ -1,7 +1,13 @@
 local ts_utils = require('nvim-treesitter.ts_utils')
 local config = require('orgmode.config')
+local query = vim.treesitter.query
 local M = {}
 
+function M.current_node()
+  local window = vim.api.nvim_get_current_win()
+  return ts_utils.get_node_at_cursor(window)
+end
+
 -- walks the tree to find a headline
 function M.find_headline(node)
   if node:type() == 'headline' then
@@ -19,7 +25,22 @@ end
 -- returns the nearest headline
 function M.closest_headline()
   vim.treesitter.get_parser(0, 'org'):parse()
-  return M.find_headline(ts_utils.get_node_at_cursor(vim.api.nvim_get_current_win()))
+  return M.find_headline(M.current_node())
+end
+
+function M.find_list(node)
+  if node:type() == 'list' then
+    return node
+  end
+  if node:type() == 'body' then
+    return nil
+  end
+  return M.find_list(node:parent())
+end
+
+function M.closest_list()
+  vim.treesitter.get_parser(0, 'org'):parse()
+  return M.find_list(M.current_node())
 end
 
 -- @param front_trim boolean

tests/plenary/ui/mappings_spec.lua

@@ -1596,4 +1596,49 @@ describe('Mappings', function()
       '*** to be refiled',
     }, vim.api.nvim_buf_get_lines(0, 0, 5, false))
   end)
+
+  it('should update the checklist cookies on a headline', function()
+    helpers.load_file_content({
+      '* Test orgmode [/]',
+      '- [ ] checkbox item',
+      '- [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('* Test orgmode [1/2]', vim.fn.getline(1))
+  end)
+
+  it('should update the checklist cookies on a parent list', function()
+    helpers.load_file_content({
+      '- Test orgmode [/]',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('- Test orgmode [1/2]', vim.fn.getline(1))
+  end)
+
+  it('should update the checklist cookies with a percentage within a headline', function()
+    helpers.load_file_content({
+      '* Test orgmode [%]',
+      '- [ ] checkbox item',
+      '- [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('* Test orgmode [50%]', vim.fn.getline(1))
+  end)
+
+  it('should update the checklist cookies with a percentage within a nested list', function()
+    helpers.load_file_content({
+      '- Test orgmode [%]',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+      '  - [ ] checkbox item',
+    })
+    vim.fn.cursor(2, 1)
+    vim.cmd([[exe "norm \<C-space>"]])
+    assert.are.same('- Test orgmode [33%]', vim.fn.getline(1))
+  end)
 end)