feat(predicates): extend has-ancestor?, add has-parent?/has-type?

1 files changed, 26 insertions, 6 deletions

diff --git a/lua/nvim-treesitter/query_predicates.lua b/lua/nvim-treesitter/query_predicates.lua
index 0c0728997..18e1401d0 100644
--- a/lua/nvim-treesitter/query_predicates.lua
+++ b/lua/nvim-treesitter/query_predicates.lua
@@ -32,22 +32,31 @@ query.add_predicate("nth?", function(match, pattern, bufnr, pred)
   return false
 end)
 
-query.add_predicate('has-ancestor?', function(match, pattern, bufnr, pred)
-  if not valid_args("has-ancestor?", pred, 2, true) then return end
+local function has_ancestor(match, pattern, bufnr, pred)
+  if not valid_args(pred[1], pred, 2) then return end
 
   local node = match[pred[2]]
-  local ancestor_type = pred[3]
+  local ancestor_types = {unpack(pred, 3)}
   if not node then return true end
 
+  local just_direct_parent = pred[1]:find('has-parent', 1, true)
+
   node = node:parent()
   while node do
-    if node:type() == ancestor_type then
+    if vim.tbl_contains(ancestor_types, node:type()) then
       return true
     end
-    node = node:parent()
+    if just_direct_parent then
+      node = nil
+    else
+      node = node:parent()
+    end
   end
   return false
-end)
+end
+query.add_predicate('has-ancestor?', has_ancestor)
+
+query.add_predicate('has-parent?', has_ancestor)
 
 query.add_predicate('is?', function(match, pattern, bufnr, pred)
   if not valid_args("is?", pred, 2) then return end
@@ -64,6 +73,17 @@ query.add_predicate('is?', function(match, pattern, bufnr, pred)
   return vim.tbl_contains(types, kind)
 end)
 
+query.add_predicate('has-type?', function(match, pattern, bufnr, pred)
+  if not valid_args(pred[1], pred, 2) then return end
+
+  local node = match[pred[2]]
+  local types = {unpack(pred, 3)}
+
+  if not node then return true end
+
+  return vim.tbl_contains(types, node:type())
+end)
+
 -- Just avoid some anoying warnings for this predicate
 query.add_predicate('set!', function() return true end)
 query.add_predicate('make-range!', function() return true end)