Treesitter Parsing (again)

lua/kulala/config/init.lua

@@ -60,6 +60,8 @@ M.defaults = {
   winbar = false,
   -- enable reading vscode rest client environment variables
   vscode_rest_client_environmentvars = false,
+  -- parse requests with tree-sitter
+  treesitter = false,
 }
 
 M.default_contenttype = {

lua/kulala/parser/init.lua

@@ -8,6 +8,7 @@ local GRAPHQL_PARSER = require("kulala.parser.graphql")
 local REQUEST_VARIABLES = require("kulala.parser.request_variables")
 local STRING_UTILS = require("kulala.utils.string")
 local PARSER_UTILS = require("kulala.parser.utils")
+local TS = require("kulala.parser.treesitter")
 local PLUGIN_TMP_DIR = FS.get_plugin_tmp_dir()
 local Scripts = require("kulala.scripts")
 local Logger = require("kulala.logger")
@@ -101,6 +102,10 @@ local function parse_body(body, variables, env)
 end
 
 M.get_document = function()
+  if CONFIG.get().treesitter then
+    return TS.get_document()
+  end
+
   local content_lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
   local content = table.concat(content_lines, "\n")
   local variables = {}
@@ -387,8 +392,15 @@ function M.parse(start_request_linenr)
     },
   }
 
-  local document_variables, requests = M.get_document()
-  local req = M.get_request_at(requests, start_request_linenr)
+  local req, document_variables
+  if CONFIG:get().treesitter then
+    document_variables = TS.get_document_variables()
+    req = TS.get_request_at(start_request_linenr)
+  else
+    local requests
+    document_variables, requests = M.get_document()
+    req = M.get_request_at(requests, start_request_linenr)
+  end
   Scripts.javascript.run("pre_request", req.scripts.pre_request)
   local env = ENV_PARSER.get_env()
 
@@ -410,7 +422,7 @@ function M.parse(start_request_linenr)
   -- We need to append the contents of the file to
   -- the body if it is a POST request,
   -- or to the URL itself if it is a GET request
-  if req.body_type == "input" then
+  if req.body_type == "input" and not CONFIG:get().treesitter then
     if req.body_path:match("%.graphql$") or req.body_path:match("%.gql$") then
       local graphql_file = io.open(req.body_path, "r")
       local graphql_query = graphql_file:read("*a")
@@ -453,15 +465,20 @@ function M.parse(start_request_linenr)
   table.insert(res.cmd, PLUGIN_TMP_DIR .. "/body.txt")
   table.insert(res.cmd, "-X")
   table.insert(res.cmd, res.method)
+
+  local is_graphql = PARSER_UTILS.contains_meta_tag(req, "graphql") or
+    PARSER_UTILS.contains_header(res.headers, "x-request-type", "GraphQL")
+  if CONFIG.get().treesitter then
+    -- treesitter parser handles graphql requests before this point
+    is_graphql = false
+  end
+
   if res.headers["content-type"] ~= nil and res.body ~= nil then
     -- check if we are a graphql query
     -- we need this here, because the user could have defined the content-type
     -- as application/json, but the body is a graphql query
     -- This can happen when the user is using http-client.env.json with DEFAULT_HEADERS.
-    if
-      PARSER_UTILS.contains_meta_tag(req, "graphql")
-      or PARSER_UTILS.contains_header(res.headers, "x-request-type", "GraphQL")
-    then
+    if is_graphql then
       local gql_json = GRAPHQL_PARSER.get_json(res.body)
       if gql_json then
         table.insert(res.cmd, "--data")
@@ -477,10 +494,7 @@ function M.parse(start_request_linenr)
     end
   else -- no content type supplied
     -- check if we are a graphql query
-    if
-      PARSER_UTILS.contains_meta_tag(req, "graphql")
-      or PARSER_UTILS.contains_header(res.headers, "x-request-type", "GraphQL")
-    then
+    if is_graphql then
       local gql_json = GRAPHQL_PARSER.get_json(res.body)
       if gql_json then
         table.insert(res.cmd, "--data")

lua/kulala/parser/inspect.lua

@@ -15,7 +15,7 @@ M.get_contents = function()
   end
   if req.body ~= nil then
     table.insert(contents, "")
-    local body_as_table = vim.split(req.body, "\r\n")
+    local body_as_table = vim.split(req.body, "\r?\n")
     for _, line in ipairs(body_as_table) do
       table.insert(contents, line)
     end

lua/kulala/parser/treesitter.lua

@@ -0,0 +1,224 @@
+local CONFIG = require('kulala.config')
+local FS = require('kulala.utils.fs')
+local STRING_UTILS = require('kulala.utils.string')
+
+local M = {}
+
+local QUERIES = {
+  section = vim.treesitter.query.parse('http', '(section (request) @request) @section'),
+  variable = vim.treesitter.query.parse('http', '(variable_declaration) @variable'),
+
+  request = vim.treesitter.query.parse('http', [[
+    (comment name: (_) value: (_)) @meta
+
+    (pre_request_script
+      ((script) @script.pre.inline
+        (#offset! @script.pre.inline 0 2 0 -2))?
+      (path)? @script.pre.file)
+
+    (request
+      header: (header)? @header
+      body: [
+        (external_body) @body.external
+        (graphql_body) @body.graphql
+      ]?) @request
+
+    (res_handler_script
+      ((script) @script.post.inline
+        (#offset! @script.post.inline 0 2 0 -2))?
+      (path)? @script.post.file)
+  ]])
+}
+
+local function text(node, metadata)
+  if not node then
+    return nil
+  end
+
+  local node_text = vim.treesitter.get_node_text(node, 0, { metadata = metadata })
+  return STRING_UTILS.trim(node_text)
+end
+
+local REQUEST_VISITORS = {
+  request = function(req, args)
+    local fields = args.fields
+    local start_line, _, end_line, _ = args.node:range()
+
+    req.url = fields.url
+    req.method = fields.method
+    req.http_version = fields.http_version
+    req.body = fields.body
+    req.start_line = start_line
+    req.block_line_count = end_line - start_line
+    req.lines_length = end_line - start_line
+
+    req.show_icon_line_number = nil
+    local show_icons = CONFIG.get().show_icons
+    if show_icons ~= nil then
+      if show_icons == 'on_request' then
+        req.show_icon_line_number = start_line + 1
+      elseif show_icons == 'above_req' then
+        req.show_icon_line_number = start_line
+      elseif show_icons == 'below_req' then
+        req.show_icon_line_number = end_line
+      end
+    end
+  end,
+
+  header = function(req, args)
+    req.headers[args.fields.name:lower()] = args.fields.value
+  end,
+
+  meta = function(req, args)
+    table.insert(req.metadata, args.fields)
+  end,
+
+  ['script.pre.inline'] = function(req, args)
+    table.insert(req.scripts.pre_request.inline, args.text)
+  end,
+
+  ['script.pre.file'] = function(req, args)
+    table.insert(req.scripts.pre_request.files, args.fields.path)
+  end,
+
+  ['script.post.inline'] = function(req, args)
+    table.insert(req.scripts.post_request.inline, args.text)
+  end,
+
+  ['script.post.file'] = function(req, args)
+    table.insert(req.scripts.post_request.files, args.fields.path)
+  end,
+
+  ['body.external'] = function(req, args)
+    local contents = FS.read_file(args.fields.path)
+    local filetype, _ = vim.filetype.match { filename = args.fields.path }
+    if filetype == 'graphql' then
+      if req.method == 'POST' then
+        req.body = string.format('{ "query": %q }', STRING_UTILS.remove_newline(contents))
+        req.headers['content-type'] = 'application/json'
+      else
+        local query = STRING_UTILS.url_encode(
+          STRING_UTILS.remove_extra_space(
+            STRING_UTILS.remove_newline(contents)
+          )
+        )
+        req.url = string.format('%s?query=%s', req.url, query)
+        req.body = nil
+      end
+    else
+      req.body = contents
+    end
+  end,
+
+  ['body.graphql'] = function(req, args)
+    local json_body = {}
+
+    for child in args.node:iter_children() do
+      if child:type() == 'graphql_data' then
+        json_body.query = text(child)
+      elseif child:type() == 'json_body' then
+        local variables_str = text(child)
+        json_body.variables = vim.fn.json_decode(variables_str)
+      end
+    end
+
+    if #json_body.query > 0 then
+      req.body = vim.fn.json_encode(json_body)
+      req.headers['content-type'] = 'application/json'
+    end
+  end,
+}
+
+local function get_root_node()
+  return vim.treesitter.get_parser(0, 'http'):parse()[1]:root()
+end
+
+local function get_fields(node)
+  local tbl = {}
+  for child, field in node:iter_children() do
+    if field then
+      tbl[field] = text(child)
+    end
+  end
+  return tbl
+end
+
+local function parse_request(section_node)
+  local req = {
+    url = '',
+    method = '',
+    http_version = '',
+    headers = {},
+    body = '',
+    metadata = {},
+    show_icon_line_number = nil,
+    redirect_response_body_to_files = {},
+    start_line = 0,
+    block_line_count = 0,
+    lines_length = 0,
+    scripts = {
+      pre_request = { inline = {}, files = {} },
+      post_request = { inline = {}, files = {} },
+    },
+  }
+
+  for i, node, metadata in QUERIES.request:iter_captures(section_node, 0) do
+    local capture = QUERIES.request.captures[i]
+
+    if REQUEST_VISITORS[capture] then
+      REQUEST_VISITORS[capture](req, {
+        node = node,
+        text = text(node, metadata[i]),
+        fields = get_fields(node),
+      })
+    end
+  end
+
+  return req
+end
+
+M.get_document_variables = function(root)
+  root = root or get_root_node()
+  local vars = {}
+
+  for _, node in QUERIES.variable:iter_captures(root, 0) do
+    local fields = get_fields(node)
+    vars[fields.name] = fields.value
+  end
+
+  return vars
+end
+
+M.get_request_at = function(line)
+  line = line or vim.fn.line('.')
+  local root = get_root_node()
+
+  for i, node in QUERIES.section:iter_captures(root, 0, line, line) do
+    if QUERIES.section.captures[i] == 'section' then
+      return parse_request(node)
+    end
+  end
+end
+
+M.get_all_requests = function(root)
+  root = root or get_root_node()
+  local requests = {}
+
+  for i, node in QUERIES.section:iter_captures(root, 0) do
+    if QUERIES.section.captures[i] == 'request' then
+      local start_line, _, end_line, _ = node:range()
+      table.insert(requests, { start_line = start_line, end_line = end_line })
+    end
+  end
+
+  return requests
+end
+
+M.get_document = function ()
+  local root = get_root_node()
+  local variables = M.get_document_variables(root)
+  local requests = M.get_all_requests(root)
+  return variables, requests
+end
+
+return M

lua/kulala/ui/init.lua

@@ -9,6 +9,8 @@ local FS = require("kulala.utils.fs")
 local DB = require("kulala.db")
 local INT_PROCESSING = require("kulala.internal_processing")
 local FORMATTER = require("kulala.formatter")
+local TS = require("kulala.parser.treesitter")
+
 local Inspect = require("kulala.parser.inspect")
 local M = {}
 
@@ -182,8 +184,14 @@ end
 
 M.open_all = function()
   INLAY.clear()
-  local _, doc = PARSER.get_document()
-  CMD.run_parser_all(doc, function(success, start, icon_linenr)
+  local requests
+  if CONFIG:get().treesitter then
+    requests = TS.get_all_requests()
+  else
+    _, requests = PARSER.get_document()
+  end
+
+  CMD.run_parser_all(requests, function(success, start, icon_linenr)
     if not success then
       if icon_linenr then
         INLAY:show_error(icon_linenr)