feature: add support for cookies

fix nvim-orgmode#81

Author: lukas-reineke

lua/orgmode/org/mappings.lua

@@ -8,6 +8,7 @@ local utils = require('orgmode.utils')
 local Files = require('orgmode.parser.files')
 local config = require('orgmode.config')
 local Help = require('orgmode.objects.help')
+local syntax = require('orgmode.org.syntax')
 
 ---@class OrgMappings
 ---@field capture Capture
@@ -142,18 +143,8 @@ function OrgMappings:global_cycle()
   return vim.cmd([[silent! norm!zx]])
 end
 
--- TODO: Add hierarchy
 function OrgMappings:toggle_checkbox()
-  local line = vim.fn.getline('.')
-  local pattern = '^(%s*[%-%+]%s*%[([%sXx%-]?)%])'
-  local checkbox, state = line:match(pattern)
-  if not checkbox then
-    return
-  end
-  local new_val = vim.trim(state) == '' and '[X]' or '[ ]'
-  checkbox = checkbox:gsub('%[[%sXx%-]?%]$', new_val)
-  local new_line = line:gsub(pattern, checkbox)
-  vim.fn.setline('.', new_line)
+  syntax.update_checkbox()
 end
 
 function OrgMappings:timestamp_up_day()
@@ -378,12 +369,12 @@ function OrgMappings:handle_return(suffix)
     return vim.cmd([[startinsert!]])
   end
 
-  if item.type == 'list' or item.type == 'listitem' then
+  if vim.tbl_contains({ 'list', 'listitem', 'cookie' }, item.type) then
     vim.cmd([[normal! ^]])
     item = Files.get_current_file():get_current_node()
   end
 
-  if item.type == 'itemtext' or item.type == 'bullet' or item.type == 'checkbox' or item.type == 'description' then
+  if vim.tbl_contains({ 'itemtext', 'bullet', 'checkbox', 'description' }, item.type) then
     local list_item = item.node:parent()
     if list_item:type() ~= 'listitem' then
       return
@@ -406,7 +397,8 @@ function OrgMappings:handle_return(suffix)
     if checkbox then
       table.insert(text_edits, {
         range = range,
-        newText = checkbox .. ' [ ] \n',
+        -- we initialize the checkbox checked, then use update_checkbox to toggle it off and update the tree
+        newText = checkbox .. ' [X] \n',
       })
     elseif plain_list then
       table.insert(text_edits, {
@@ -442,6 +434,7 @@ function OrgMappings:handle_return(suffix)
       vim.lsp.util.apply_text_edits(text_edits, 0)
 
       vim.fn.cursor(end_row + 2 + (add_empty_line and 1 or 0), 0) -- +1 for 0 index and +1 for next line
+      syntax.update_checkbox(ts_utils.get_next_node(list_item))
       vim.cmd([[startinsert!]])
     end
   end

lua/orgmode/org/syntax.lua

@@ -1,5 +1,10 @@
 local Files = require('orgmode.parser.files')
 local config = require('orgmode.config')
+local utils = require('orgmode.utils')
+local ts_utils = require('nvim-treesitter.ts_utils')
+
+local checkbox_query = vim.treesitter.parse_query('org', '(list (listitem (checkbox) @checkbox))')
+local headline_cookie_query = vim.treesitter.parse_query('org', '(headline (cookie) @cookie)')
 
 local function load_code_blocks()
   local file = vim.api.nvim_buf_get_name(0)
@@ -36,7 +41,115 @@ local function add_todo_keywords_to_spellgood()
   end
 end
 
+local function _get_cookie_checked_and_total(parent)
+  local parent_type = parent:type()
+  local start_row, _, end_row, _ = parent:range()
+  local checked, total = 0, 0
+  for _, checkbox in checkbox_query:iter_captures(parent, 0, start_row, end_row + 1) do
+    local closest_parent = utils.get_closest_parent_of_type(checkbox, parent_type)
+    if closest_parent and closest_parent == parent then -- only count direct children
+      local checkbox_text = vim.treesitter.get_node_text(checkbox, 0)
+      if checkbox_text:match('%[[x|X]%]') then
+        checked = checked + 1
+      end
+      total = total + 1
+    end
+  end
+
+  return checked, total
+end
+
+local function _update_checkbox_text(checkbox, checked_children, total_children)
+  local checkbox_text
+  if total_children == nil then -- if the function is called without child information, we toggle the current value
+    checkbox_text = vim.treesitter.get_node_text(checkbox, 0)
+    if checkbox_text:match('%[[xX]%]') then
+      checkbox_text = '[ ]'
+    else
+      checkbox_text = '[X]'
+    end
+  else
+    checkbox_text = '[ ]'
+    if checked_children == total_children then
+      checkbox_text = '[x]'
+    elseif checked_children > 0 then
+      checkbox_text = '[-]'
+    end
+  end
+
+  utils.update_node_text(checkbox, { checkbox_text })
+end
+
+local function _update_cookie_text(cookie, checked_children, total_children)
+  local cookie_text = vim.treesitter.get_node_text(cookie, 0)
+
+  if total_children == nil then
+    checked_children, total_children = 0, 0
+  end
+
+  local new_cookie
+  if cookie_text:find('/') then
+    new_cookie = string.format('[%d/%d]', checked_children, total_children)
+  else
+    if total_children > 0 then
+      new_cookie = string.format('[%d%%%%]', (100 * checked_children) / total_children)
+    else
+      new_cookie = '[0%%%]'
+    end
+  end
+  cookie_text = cookie_text:gsub('%[.*%]', new_cookie)
+  utils.update_node_text(cookie, { cookie_text })
+end
+
+local function update_checkbox(node, checked_children, total_children)
+  if not node then
+    node = utils.get_closest_parent_of_type(ts_utils.get_node_at_cursor(0), 'listitem')
+    if not node then
+      return
+    end
+  end
+
+  local checkbox
+  local cookie
+  for child in node:iter_children() do
+    if child:type() == 'checkbox' then
+      checkbox = child
+    elseif child:type() == 'itemtext' then
+      local c_child = child:named_child(0)
+      if c_child and c_child:type() == 'cookie' then
+        cookie = c_child
+      end
+    end
+  end
+
+  if checkbox then
+    _update_checkbox_text(checkbox, checked_children, total_children)
+  end
+
+  if cookie then
+    _update_cookie_text(cookie, checked_children, total_children)
+  end
+
+  local listitem_parent = utils.get_closest_parent_of_type(node:parent(), 'listitem')
+  if listitem_parent then
+    local list_parent = utils.get_closest_parent_of_type(node, 'list')
+    local checked, total = _get_cookie_checked_and_total(list_parent)
+    return update_checkbox(listitem_parent, checked, total)
+  end
+
+  local section = utils.get_closest_parent_of_type(node:parent(), 'section')
+  if section then
+    local list_parent = utils.get_closest_parent_of_type(node, 'list')
+    local checked, total = _get_cookie_checked_and_total(list_parent)
+    local start_row, _, end_row, _ = section:range()
+    for _, headline_cookie in headline_cookie_query:iter_captures(section, 0, start_row, end_row + 1) do
+      _update_cookie_text(headline_cookie, checked, total)
+    end
+  end
+end
+
 return {
   load_code_blocks = load_code_blocks,
   add_todo_keywords_to_spellgood = add_todo_keywords_to_spellgood,
+  update_checkbox = update_checkbox,
 }

lua/orgmode/utils/init.lua

@@ -321,6 +321,13 @@ function utils.get_node_text(node, content)
   end
 end
 
+---@param node userdata
+---@param text string[]
+function utils.update_node_text(node, text)
+  local start_row, start_col, end_row, end_col = node:range()
+  vim.api.nvim_buf_set_text(0, start_row, start_col, end_row, end_col, text)
+end
+
 ---@param node table
 ---@param type string
 ---@return table