feat: support headline's cookie from TODOs

lua/orgmode/files/headline.lua

@@ -357,12 +357,14 @@ end
 function Headline:set_todo(keyword)
   local todo, node = self:get_todo()
   if todo then
-    return self:_set_node_text(node, keyword)
+    self:_set_node_text(node, keyword)
+    return self:update_parent_cookie()
   end
 
   local stars = self:_get_child_node('stars')
   local _, level = stars:end_()
-  return self:_set_node_text(stars, ('%s %s'):format(('*'):rep(level), keyword))
+  self:_set_node_text(stars, ('%s %s'):format(('*'):rep(level), keyword))
+  return self:update_parent_cookie()
 end
 
 memoize('get_todo')
@@ -890,36 +892,74 @@ function Headline:get_cookie()
   return self:_parse_title_part('%[%d?%d?%d?%%%]')
 end
 
+function Headline:_set_cookie(cookie, num, denum)
+  -- Update the cookie
+  local new_cookie_val
+  if self.file:get_node_text(cookie):find('%%') then
+    new_cookie_val = ('[%d%%]'):format((num / denum) * 100)
+  else
+    new_cookie_val = ('[%d/%d]'):format(num, denum)
+  end
+  return self:_set_node_text(cookie, new_cookie_val)
+end
+
 function Headline:update_cookie()
+  -- Update cookie state from a check box state change
+
+  -- Return early if the headline doesn't have a cookie
+  local cookie = self:get_cookie()
+  if not cookie then
+    return self
+  end
+
   local section = self:node():parent()
   if not section then
     return self
   end
 
-  -- Go through all the lists in this headline and gather checked_boxes
-  local num_boxes, num_checked_boxes = 0, 0
+  -- Count checked boxes from all lists
+  local num_checked_boxes, num_boxes = 0, 0
   local body = section:field('body')[1]
-  for node in body:iter_children() do
-    if node:type() == 'list' then
-      local boxes = self:child_checkboxes(node)
-      num_boxes = num_boxes + #boxes
-      local checked_boxes = vim.tbl_filter(function(box)
-        return box:match('%[%w%]')
-      end, boxes)
-      num_checked_boxes = num_checked_boxes + #checked_boxes
+  if body then
+    for node in body:iter_children() do
+      if node:type() == 'list' then
+        local boxes = self:child_checkboxes(node)
+        num_boxes = num_boxes + #boxes
+        local checked_boxes = vim.tbl_filter(function(box)
+          return box:match('%[%w%]')
+        end, boxes)
+        num_checked_boxes = num_checked_boxes + #checked_boxes
+      end
     end
   end
 
-  -- Update the cookie
+  -- Set the cookie
+  return self:_set_cookie(cookie, num_checked_boxes, num_boxes)
+end
+
+function Headline:update_todo_cookie()
+  -- Update cookie state from a TODO state change
+
+  -- Return early if the headline doesn't have a cookie
   local cookie = self:get_cookie()
-  if cookie then
-    local new_cookie_val
-    if self.file:get_node_text(cookie):find('%%') then
-      new_cookie_val = ('[%d%%]'):format((num_checked_boxes / num_boxes) * 100)
-    else
-      new_cookie_val = ('[%d/%d]'):format(num_checked_boxes, num_boxes)
-    end
-    return self:_set_node_text(cookie, new_cookie_val)
+  if not cookie then
+    return self
+  end
+
+  -- Count done children headlines
+  local children = self:get_child_headlines()
+  local dones = vim.tbl_filter(function(h)
+    return h:is_done()
+  end, children)
+
+  -- Set the cookie
+  return self:_set_cookie(cookie, #dones, #children)
+end
+
+function Headline:update_parent_cookie()
+  local parent = self:get_parent_headline()
+  if parent and parent.headline then
+    parent:update_todo_cookie()
   end
   return self
 end

tests/plenary/ui/mappings/todo_spec.lua

@@ -408,4 +408,26 @@ describe('Todo mappings', function()
       '  :END:',
     }, vim.api.nvim_buf_get_lines(0, 2, 11, false))
   end)
+
+  it('should update headline cookies when children todo state changes', function()
+    helpers.create_file({
+      '* Test orgmode [/]',
+      '** TODO item',
+      '** TODO item',
+      '** TODO item',
+      '** TODO item',
+    })
+    vim.fn.cursor(4, 1)
+    local now = Date.now()
+    -- Changing to DONE and adding closed date
+    vim.cmd([[norm citd]])
+    assert.are.same({
+      '* Test orgmode [1/4]',
+      '** TODO item',
+      '** TODO item',
+      '** DONE item',
+      '   CLOSED: [' .. now:to_string() .. ']',
+      '** TODO item',
+    }, vim.api.nvim_buf_get_lines(0, 0, 6, false))
+  end)
 end)