fix: add treesitter support for JS / TS / TSX generator functions (#289)

* fix: add JS as TS generator functions

* feat: add TSX generator functions

lua/aerial/backends/treesitter/extensions.lua

@@ -309,6 +309,9 @@ M.typescript = {
   postprocess = function(bufnr, item, match)
     local value_node = (utils.get_at_path(match, "var_type") or {}).node
     if value_node then
+      if value_node:type() == "generator_function" then
+        item.kind = "Function"
+      end
       if value_node:type() == "arrow_function" then
         item.kind = "Function"
       end

queries/javascript/aerial.scm

@@ -8,6 +8,11 @@
   (#set! "kind" "Function")
   ) @type
 
+(generator_function_declaration
+  name: (identifier) @name
+  (#set! "kind" "Function")
+  ) @type
+
 (method_definition
   name: (property_identifier) @name
   (#set! "kind" "Method")
@@ -23,7 +28,7 @@
 (lexical_declaration
   (variable_declarator
     name: (identifier) @name
-    value: [(arrow_function) (function)] @type
+    value: [(arrow_function) (function) (generator_function)] @type
   )
   (#set! "kind" "Function")
   ) @start

queries/tsx/aerial.scm

@@ -8,6 +8,11 @@
   (#set! "kind" "Function")
   ) @type
 
+(generator_function_declaration
+  name: (identifier) @name
+  (#set! "kind" "Function")
+  ) @type
+
 (interface_declaration
   name: (type_identifier) @name
   (#set! "kind" "Interface")

queries/typescript/aerial.scm

@@ -8,6 +8,11 @@
   (#set! "kind" "Function")
   ) @type
 
+(generator_function_declaration
+  name: (identifier) @name
+  (#set! "kind" "Function")
+  ) @type
+
 (interface_declaration
   name: (type_identifier) @name
   (#set! "kind" "Interface")

tests/treesitter/js_spec.lua

@@ -162,6 +162,24 @@ describe("treesitter js", function()
         end_lnum = 23,
         end_col = 28,
       },
+      {
+        kind = "Function",
+        name = "g",
+        level = 0,
+        lnum = 25,
+        col = 0,
+        end_lnum = 25,
+        end_col = 16,
+      },
+      {
+        kind = "Function",
+        name = "g_1",
+        level = 0,
+        lnum = 26,
+        col = 0,
+        end_lnum = 26,
+        end_col = 28,
+      },
     })
   end)
 end)

tests/treesitter/js_test.js

@@ -21,3 +21,6 @@ describe("UnitTest", () => {
 
 const fn_2 = () => {};
 const fn_3 = function () {};
+
+function* g() {}
+const g_1 = function* () {};

tests/treesitter/ts_spec.lua

@@ -198,6 +198,24 @@ describe("treesitter ts", function()
           },
         },
       },
+      {
+        kind = "Function",
+        name = "g",
+        level = 0,
+        lnum = 33,
+        col = 0,
+        end_lnum = 33,
+        end_col = 16,
+      },
+      {
+        kind = "Function",
+        name = "g_1",
+        level = 0,
+        lnum = 34,
+        col = 0,
+        end_lnum = 34,
+        end_col = 28,
+      },
     })
   end)
 end)

tests/treesitter/ts_test.ts

@@ -29,3 +29,6 @@ describe("UnitTest", () => {
     test.each([])("runs multiple times", () => {});
   });
 });
+
+function* g() {}
+const g_1 = function* () {};

tests/treesitter/tsx_spec.lua

@@ -198,6 +198,24 @@ describe("treesitter tsx", function()
           },
         },
       },
+      {
+        kind = "Function",
+        name = "g",
+        level = 0,
+        lnum = 35,
+        col = 0,
+        end_lnum = 35,
+        end_col = 16,
+      },
+      {
+        kind = "Function",
+        name = "g_1",
+        level = 0,
+        lnum = 36,
+        col = 0,
+        end_lnum = 36,
+        end_col = 28,
+      },
     })
   end)
 end)

tests/treesitter/tsx_test.tsx

@@ -31,3 +31,6 @@ describe("UnitTest", () => {
     test.each([])("runs multiple times", () => {});
   });
 });
+
+function* g() {}
+const g_1 = function* () {};