Handle nested lists with cookies

Author: TravonteD

lua/orgmode/org/mappings.lua

@@ -165,7 +165,7 @@ function OrgMappings:toggle_checkbox()
   vim.fn.setline('.', new_line)
   local headline = Headline:new(tree_utils.closest_headline())
   local list = List:new(tree_utils.closest_list())
-  headline:update_cookie(list)
+  list:update_parent_cookie()
 end
 
 function OrgMappings:timestamp_up_day()

lua/orgmode/treesitter/headline.lua

@@ -124,20 +124,6 @@ function Headline:cookie()
   return self:parse('%[%d?/%d?%]')
 end
 
-function Headline:update_cookie(list)
-  local checkbox_status = self:cookie()
-  if not checkbox_status then
-    return
-  end
-
-  local checkboxes = list:checkboxes()
-  local checked_boxes = vim.tbl_filter(function(box)
-    return box:match('%[%w%]')
-  end, checkboxes)
-  local new_status = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
-  tree_utils.set_node_text(checkbox_status, new_status)
-end
-
 -- @return tsnode, string
 function Headline:parse(pattern)
   local match = ''

lua/orgmode/treesitter/list.lua

@@ -1,6 +1,7 @@
 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 = {}
 
@@ -11,6 +12,40 @@ function List:new(list_node)
   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)
+    if query.get_node_text(cookie_node, 0):match('%[%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 = ('[%d/%d]'):format(#checked_boxes, #checkboxes)
+  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)

tests/plenary/ui/mappings_spec.lua

@@ -1607,4 +1607,15 @@ describe('Mappings', function()
     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)
 end)