refactor: Use custom predicate for todo keyword highlights

Author: kristijanhusak

lua/orgmode/colors/todo_highlighter.lua

@@ -8,8 +8,11 @@ local function add_todo_keyword_highlights()
   if not query_files or #query_files == 0 then
     return
   end
-  local todo_keywords = config:get_todo_keywords()
   local faces = highlights.parse_todo_keyword_faces()
+  if not faces or vim.tbl_isempty(faces) then
+    return
+  end
+
   local all_lines = {}
   for i, _ in pairs(query_files) do
     if i ~= #query_files then
@@ -31,20 +34,6 @@ local function add_todo_keyword_highlights()
           if err then
             return
           end
-          local todo_type = table.concat(
-            vim.tbl_map(function(word)
-              return string.format('"%s"', word)
-            end, todo_keywords.TODO),
-            ' '
-          )
-          local done_type = table.concat(
-            vim.tbl_map(function(word)
-              return string.format('"%s"', word)
-            end, todo_keywords.DONE),
-            ' '
-          )
-          table.insert(lines, string.format([[(item . (expr) @OrgTODO @nospell (#any-of? @OrgTODO %s))]], todo_type))
-          table.insert(lines, string.format([[(item . (expr) @OrgDONE @nospell (#any-of? @OrgDONE %s))]], done_type))
           for face_name, face_hl in pairs(faces) do
             table.insert(
               lines,

lua/orgmode/config/init.lua

@@ -343,6 +343,20 @@ function Config:get_inheritable_tags(headline)
   end, headline.tags)
 end
 
+function Config:setup_ts_predicates()
+  local todo_keywords = self:get_todo_keywords().KEYS
+
+  vim.treesitter.query.add_predicate('org-is-todo-keyword?', function(match, _, source, predicate)
+    local node = match[predicate[2]]
+    if node then
+      local text = vim.treesitter.get_node_text(node, source)
+      return todo_keywords[text] and todo_keywords[text].type == predicate[3] or false
+    end
+
+    return false
+  end, true)
+end
+
 function Config:ts_highlights_enabled()
   if self.ts_hl_enabled ~= nil then
     return self.ts_hl_enabled

lua/orgmode/init.lua

@@ -108,6 +108,7 @@ local function setup(opts)
   instance = Org:new()
   check_ts_grammar()
   local config = require('orgmode.config'):extend(opts)
+  config:setup_ts_predicates()
   vim.defer_fn(function()
     if config.notifications.enabled and #vim.api.nvim_list_uis() > 0 then
       require('orgmode.parser.files').load(vim.schedule_wrap(function()

queries/org/highlights.scm

@@ -9,6 +9,8 @@
 (headline (stars) @stars (#eq? @stars "*******")) @OrgTSHeadlineLevel7
 (headline (stars) @stars (#eq? @stars "********")) @OrgTSHeadlineLevel8
 (headline (item) @spell)
+(item . (expr) @OrgTODO @nospell (#org-is-todo-keyword? @OrgTODO "TODO"))
+(item . (expr) @OrgDONE @nospell (#org-is-todo-keyword? @OrgDONE "DONE"))
 (list (listitem (paragraph) @spell))
 (body (paragraph) @spell)
 (bullet) @OrgTSBullet