refactor: use treesitter to add properties to a headline

Author: kristijanhusak

lua/orgmode/clock/init.lua

@@ -1,4 +1,5 @@
 local Files = require('orgmode.parser.files')
+local tree_utils = require('orgmode.treesitter')
 local Duration = require('orgmode.objects.duration')
 local utils = require('orgmode.utils')
 local Promise = require('orgmode.utils.promise')
@@ -75,18 +76,15 @@ function Clock:org_clock_goto()
 end
 
 function Clock:org_set_effort()
-  local item = Files.get_closest_headline()
-  if not item then
-    return
-  end
+  local item = tree_utils.closest_headline()
   -- TODO: Add Effort_ALL property as autocompletion
   local current_effort = item:get_property('Effort')
-  local effort = vim.fn.OrgmodeInput('Effort: ', current_effort or '')
+  local effort = vim.fn.OrgmodeInput('Effort: ', current_effort and current_effort.value or '')
   local duration = Duration.parse(effort)
   if duration == nil then
     return utils.echo_error('Invalid duration format: ' .. effort)
   end
-  item:add_properties({ Effort = effort })
+  item:set_property('Effort', effort)
 end
 
 function Clock:get_statusline()

lua/orgmode/config/init.lua

@@ -353,12 +353,28 @@ function Config:respect_blank_before_new_entry(content, option, prepend_content)
   return content
 end
 
+---@param amount number
+---@return string
 function Config:get_indent(amount)
   if self.opts.org_indent_mode == 'indent' then
     return string.rep(' ', amount)
   end
   return ''
 end
 
+---@param content table|string
+---@param amount number
+function Config:apply_indent(content, amount)
+  local indent = self:get_indent(amount)
+  if type(content) ~= 'table' then
+    return indent .. content
+  end
+
+  for i, line in ipairs(content) do
+    content[i] = indent .. line
+  end
+  return content
+end
+
 instance = Config:new()
 return instance

lua/orgmode/org/mappings.lua

@@ -49,20 +49,14 @@ function OrgMappings:archive()
   Files.reload(
     archive_location,
     vim.schedule_wrap(function()
-      Files.update_file(archive_location, function(archive_file)
-        -- Jump to last headline in archive file
-        vim.cmd('normal! G')
-        vim.fn.search([[^\*\+\s\+]], 'b')
-        local last_item = archive_file:get_closest_headline()
-        if not last_item then
-          return
+      Files.update_file(archive_location, function()
+        local archived_headline = ts_org.find_headline(item.title, true)
+        if archived_headline then
+          archived_headline:set_property('ARCHIVE_TIME', Date.now():to_string())
+          archived_headline:set_property('ARCHIVE_FILE', file.filename)
+          archived_headline:set_property('ARCHIVE_CATEGORY', item.category)
+          archived_headline:set_property('ARCHIVE_TODO', item.todo_keyword.value)
         end
-        last_item:add_properties({
-          ARCHIVE_TIME = Date.now():to_string(),
-          ARCHIVE_FILE = file.filename,
-          ARCHIVE_CATEGORY = item.category,
-          ARCHIVE_TODO = item.todo_keyword.value,
-        })
       end)
     end)
   )
@@ -444,8 +438,9 @@ function OrgMappings:_todo_change_state(direction)
   end
 
   self:_change_todo_state('reset')
-  local state_change =
-    string.format('%s- State "%s" from "%s" [%s]', indent, item.todo_keyword.value, old_state, Date.now():to_string())
+  local state_change = {
+    string.format('%s- State "%s" from "%s" [%s]', indent, item.todo_keyword.value, old_state, Date.now():to_string()),
+  }
 
   dispatchEvent()
   return Promise.resolve()
@@ -459,20 +454,12 @@ function OrgMappings:_todo_change_state(direction)
       end)
     end)
     :next(function(note)
-      local data = item:add_properties({ LAST_REPEAT = '[' .. Date.now():to_string() .. ']' })
+      headline:set_property('LAST_REPEAT', Date.now():to_wrapped_string(false))
       if not note then
         return
       end
-
-      if data.is_new then
-        vim.fn.append(data.end_line, note)
-        return
-      end
-      item = Files.get_closest_headline()
-
-      if item.properties.valid then
-        vim.fn.append(item.properties.range.end_line, note)
-      end
+      local properties_end_line = headline:properties():end_()
+      vim.api.nvim_buf_set_lines(0, properties_end_line, properties_end_line, false, note)
     end)
 end
 

lua/orgmode/parser/section.lua

@@ -321,6 +321,7 @@ function Section:get_repeater_dates()
   end, self.dates)
 end
 
+---@deprecated use treesitter.headline.set_property
 ---@param properties table
 ---@return table
 function Section:add_properties(properties)

lua/orgmode/treesitter/headline.lua

@@ -17,11 +17,31 @@ function Headline:new(headline_node)
   return data
 end
 
+---@param cursor? Table Cursor position tuple {row, col}
+---@return Headline|nil
+function Headline.from_cursor(cursor)
+  local ts_headline = tree_utils.closest_headline(cursor)
+  if not ts_headline then
+    return nil
+  end
+  return Headline:new(ts_headline)
+end
+
 ---@return userdata stars node
 function Headline:stars()
   return self.headline:field('stars')[1]
 end
 
+function Headline:refresh()
+  tree_utils.parse()
+  local start_row, start_col = self.headline:start()
+  local updated_headline = Headline.from_cursor({ start_row + 1, start_col })
+  if updated_headline then
+    self.headline = updated_headline.headline
+  end
+  return self
+end
+
 ---@return number
 function Headline:level()
   local stars = self:stars()
@@ -150,6 +170,15 @@ function Headline:todo()
   end
 end
 
+function Headline:title()
+  local title = query.get_node_text(self:item(), 0) or ''
+  local todo, word = self:todo()
+  if todo then
+    title = title:gsub('^' .. vim.pesc(word) .. '%s*', '')
+  end
+  return title
+end
+
 ---@return userdata
 function Headline:plan()
   local section = self.headline:parent()
@@ -160,6 +189,75 @@ function Headline:plan()
   end
 end
 
+---@return userdata
+function Headline:properties()
+  local section = self.headline:parent()
+  for _, node in ipairs(ts_utils.get_named_children(section)) do
+    if node:type() == 'property_drawer' then
+      return node
+    end
+  end
+end
+
+---@param name string
+---@param value string
+function Headline:set_property(name, value)
+  local properties = self:properties()
+  if not properties then
+    local append_line = self:get_append_line()
+    local property_drawer = self:_apply_indent({ ':PROPERTIES:', ':END:' })
+    vim.api.nvim_buf_set_lines(0, append_line, append_line, false, property_drawer)
+    tree_utils.parse()
+    properties = self:refresh():properties()
+  end
+
+  local property = (':%s: %s'):format(name, value)
+  local existing_property = self:get_property(name)
+  if existing_property then
+    tree_utils.set_node_text(existing_property.node, property)
+    return self:refresh()
+  end
+  local property_end = properties:end_()
+  vim.api.nvim_buf_set_lines(0, property_end - 1, property_end - 1, false, { self:_apply_indent(property) })
+  return self:refresh()
+end
+
+---@param property_name string
+---@return table|nil
+function Headline:get_property(property_name)
+  local properties = self:properties()
+  if not properties then
+    return nil
+  end
+
+  for _, node in ipairs(ts_utils.get_named_children(properties)) do
+    local name = node:field('name')[1]
+    local value = node:field('value')[1]
+    if name and query.get_node_text(name, 0):lower() == property_name:lower() then
+      return {
+        node = node,
+        name = name,
+        value = value and query.get_node_text(value, 0),
+      }
+    end
+  end
+end
+
+---Return the line number where content can be appended
+---
+---@return number
+function Headline:get_append_line()
+  local properties = self:properties()
+  if properties then
+    return properties:end_()
+  end
+  local plan = self:plan()
+  if plan then
+    return plan:end_()
+  end
+  return self.headline:end_()
+end
+
 ---@return Table<string, userdata>
 function Headline:dates()
   local plan = self:plan()
@@ -341,4 +439,9 @@ function Headline:_remove_date(type)
   end
 end
 
+---@param text table|string
+function Headline:_apply_indent(text)
+  return config:apply_indent(text, self:level() + 1)
+end
+
 return Headline

lua/orgmode/treesitter/init.lua

@@ -5,11 +5,11 @@ local Listitem = require('orgmode.treesitter.listitem')
 ---@param cursor? Table Cursor position tuple {row, col}
 ---@return Headline
 local function closest_headline(cursor)
-  local ts_headline = tree_utils.closest_headline(cursor)
+  local ts_headline = Headline.from_cursor(cursor)
   if not ts_headline then
     error('Unable to locate closest headline')
   end
-  return Headline:new(ts_headline)
+  return ts_headline
 end
 
 local function listitem()
@@ -21,7 +21,33 @@ local function listitem()
   return nil
 end
 
+---@param title string
+---@param exact? boolean
+local function find_headline(title, exact)
+  local trees = vim.treesitter.get_parser(0, 'org', {}):parse()
+  if #trees == 0 then
+    return nil
+  end
+  local root = trees[1]:root()
+  local ts_query = tree_utils.parse_query('(section (headline) @headline)')
+  for _, match, _ in ts_query:iter_matches(root) do
+    -- local items = {}
+    for _, matched_node in pairs(match) do
+      local headline = Headline:new(matched_node)
+      local pattern = '^' .. vim.pesc(title:lower())
+      if exact then
+        pattern = pattern .. '$'
+      end
+
+      if headline:title():lower():match(pattern) then
+        return headline
+      end
+    end
+  end
+end
+
 return {
   closest_headline = closest_headline,
   listitem = listitem,
+  find_headline = find_headline,
 }

lua/orgmode/utils/treesitter.lua

@@ -1,12 +1,26 @@
 local ts_utils = require('nvim-treesitter.ts_utils')
 local parsers = require('nvim-treesitter.parsers')
 local M = {}
+local query_cache = {}
 
 function M.current_node()
   local window = vim.api.nvim_get_current_win()
   return ts_utils.get_node_at_cursor(window)
 end
 
+function M.parse(bufnr)
+  return vim.treesitter.get_parser(bufnr or 0, 'org', {}):parse()
+end
+
+function M.parse_query(query)
+  local ts_query = query_cache[query]
+  if not ts_query then
+    ts_query = vim.treesitter.query.parse_query('org', query)
+    query_cache[query] = ts_query
+  end
+  return query_cache[query]
+end
+
 ---This is a full copy of nvim_treesiter get_node_at_cursor with support for custom cursor position
 ---@param cursor? Table Cursor position tuple {row, col}
 ---@param winnr? number