new implementation

lua/ts_context_commentstring/internal.lua

@@ -51,10 +51,8 @@ end
 function M.calculate_commentstring(args)
   args = args or {}
   local key = args.key or '__default'
-  local location = args.location or nil
 
-  local node, language_tree =
-    utils.get_node_at_cursor_start_of_line(vim.tbl_keys(config.get_languages_config()), location)
+  local node, language_tree = utils.get_node_at_cursor_start_of_line()
 
   if not node and not language_tree then
     return nil
@@ -74,7 +72,7 @@ function M.calculate_commentstring(args)
   -- Use ts node commentstring if no language is configured
   if not language_config then
     local cs = M.check_ts_node(language_tree)
-    if cs ~= '' then
+    if cs then
       return cs
     end
   end
@@ -83,26 +81,10 @@ function M.calculate_commentstring(args)
 end
 
 function M.check_ts_node(language_tree)
-  local ts_cs = nil
-  local traverse
-
-  traverse = function(lang_tree)
-    local lang = lang_tree:lang()
-    local filetypes = vim.treesitter.language.get_filetypes(lang)
-    for _, ft in ipairs(filetypes) do
-      local cur_cs = vim.filetype.get_option(ft, 'commentstring')
-      if type(cur_cs) == 'string' and cur_cs ~= '' then
-        ts_cs = cur_cs
-      end
-    end
-
-    for _, child_lang_tree in pairs(lang_tree:children()) do
-      traverse(child_lang_tree)
-    end
+  local cs = vim.filetype.get_option(language_tree:lang(), 'commentstring')
+  if type(cs) == 'string' and cs ~= '' then
+    return cs
   end
-  traverse(language_tree)
-
-  return ts_cs
 end
 
 ---Update the `commentstring` setting based on the current location of the

lua/ts_context_commentstring/utils.lua

@@ -74,20 +74,14 @@ end
 ---Returns the node at the cursor's line and the language tree for that
 ---injection.
 ---
----@param only_languages string[] List of languages to filter for, all
----  other languages will be ignored.
----@param location? ts_context_commentstring.Location location Line, column
----  where to start traversing the tree. Defaults to cursor start of line.
----  This usually makes the most sense when commenting the whole line.
----
 ---@return table|nil node, table|nil language_tree Node and language tree for the
 ---  location
-function M.get_node_at_cursor_start_of_line(only_languages, location)
+function M.get_node_at_cursor_start_of_line()
   if not M.is_treesitter_active() then
     return
   end
 
-  location = location or M.get_cursor_line_non_whitespace_col_location()
+  local location = M.get_cursor_line_non_whitespace_col_location()
   local range = {
     location[1],
     location[2],
@@ -99,7 +93,9 @@ function M.get_node_at_cursor_start_of_line(only_languages, location)
   local language_tree = vim.treesitter.get_parser()
   -- Get the smallest supported language's tree with nodes inside the given range
   language_tree:for_each_tree(function(_, ltree)
-    if ltree:contains(range) and vim.tbl_contains(only_languages, ltree:lang()) then
+    -- always exclude the comment language as it's returning invalid
+    -- commentstring
+    if ltree:contains(range) and ltree:lang() ~= 'comment' then
       language_tree = ltree
     end
   end)