feat(rust): test_fn now can expand rstest attr if the test mod has already use rstest

Author: TwIStOy

README.md

@@ -217,7 +217,7 @@ tmss! -> absl::flat_hash_map<std::string, std::string>
 |   `.ts`   | Switches indent's coding style between `CamelCase` and `snake_case`. |      `indent`      |
 |   `.sc`   | Wraps with `static_cast<>(?)`.                                       |     `any_expr`     |
 | `.single` | Wraps with `ranges::views::single(?)`.                               |     `any_expr`     |
-| `.await`  | Expands to `co_await ?`.                                              |     `any_expr`     |
+| `.await`  | Expands to `co_await ?`.                                             |     `any_expr`     |
 |   `.in`   | Expands to `if (...find)` statements.                                |     `any_expr`     |
 
 </details>

lua/luasnip-snippets/config.lua

@@ -17,12 +17,16 @@
 ---@class LSSnippets.Config.Snippet.Cpp
 ---@field quick_type? LSSnippets.Config.Snippet.Cpp.QuickType
 
+---@class LSSnippets.Config.Snippet.Rust
+---@field rstest_support? boolean
+
 ---@alias LSSnippets.Config.Snippet.DisableSnippets string[]
 ---@alias LSSnippets.SupportLangs 'cpp'|'dart'|'lua'|'rust'|'nix'|'typescript'|'*'
 
 ---@class LSSnippets.Config.Snippet
 ---@field lua? LSSnippets.Config.Snippet.Lua
 ---@field cpp? LSSnippets.Config.Snippet.Cpp
+---@field rust? LSSnippets.Config.Snippet.Rust
 
 ---@class LSSnippets.Config
 ---@field copyright_header? string
@@ -35,7 +39,7 @@ local config = {}
 ---@param opts? LSSnippets.Config
 local function setup(opts)
   opts = opts or {}
-  config = vim.tbl_extend("force", config, opts)
+  config = vim.tbl_deep_extend("force", config, opts)
 end
 
 ---@return any

lua/luasnip-snippets/snippets/rust/test_fn.lua

@@ -1,4 +1,6 @@
 local UtilsTS = require("luasnip-snippets.utils.treesitter")
+local Config = require("luasnip-snippets.config")
+local UtilsTbl = require("luasnip-snippets.utils.tbl")
 local ls = require("luasnip")
 local d = ls.dynamic_node
 local sn = ls.snippet_node
@@ -8,6 +10,66 @@ local fmta = require("luasnip.extras.fmt").fmta
 local i = require("luasnip-snippets.nodes").insert_node
 local c = require("luasnip-snippets.nodes").choice_node
 
+---@param left string[]
+---@param right string[]
+---@param sep string
+---@return string[]
+local function dot_concat(left, right, sep)
+  local ret = {}
+
+  for _, l in ipairs(left) do
+    for _, r in ipairs(right) do
+      ret[#ret + 1] = l .. sep .. r
+    end
+  end
+
+  return ret
+end
+
+---@param node TSNode?
+---@return string[]
+local function flat_scoped_use_list(source, node)
+  if node == nil then
+    return {}
+  end
+  if node:type() ~= "scoped_use_list" then
+    return {
+      vim.treesitter.get_node_text(node, source),
+    }
+  end
+
+  local path_nodes = node:field("path")
+  if #path_nodes == 0 then
+    return {}
+  end
+
+  local paths = {}
+  for _, path_node in ipairs(path_nodes) do
+    vim.list_extend(paths, flat_scoped_use_list(source, path_node))
+  end
+
+  local items = {}
+  local name_nodes = node:field("name")
+  for _, name_node in ipairs(name_nodes) do
+    vim.list_extend(items, flat_scoped_use_list(source, name_node))
+  end
+  local list_nodes = node:field("list")
+  local allow_list = {
+    scoped_use_list = 1,
+    use_wildcard = 1,
+    identifier = 1,
+  }
+  for _, list_node in ipairs(list_nodes) do
+    for child in list_node:iter_children() do
+      if allow_list[child:type()] == 1 then
+        vim.list_extend(items, flat_scoped_use_list(source, child))
+      end
+    end
+  end
+
+  return dot_concat(paths, items, "::")
+end
+
 local function inject_expanding_environment(_, _, match, captures)
   local row, col = unpack(vim.api.nvim_win_get_cursor(0))
   local buf = vim.api.nvim_get_current_buf()
@@ -73,6 +135,31 @@ local function inject_expanding_environment(_, _, match, captures)
           prev = prev:prev_sibling()
         end
         ret.env_override["ATTRIBUTES_ITEMS"] = attributes
+
+        if Config.get("snippet.rust.rstest_support") == true then
+          -- check if this mod contains `use rstest::rstest;`
+          local use_list = {}
+          for _, body in ipairs(mod_item:field("body")) do
+            for child in body:iter_children() do
+              if child:type() == "use_declaration" then
+                local nodes = child:field("argument")
+                for _, use_node in ipairs(nodes) do
+                  local node_type = use_node:type()
+                  if node_type == "scoped_use_list" then
+                    vim.list_extend(
+                      use_list,
+                      flat_scoped_use_list(source, use_node)
+                    )
+                  elseif node_type == "scoped_identifier" then
+                    use_list[#use_list + 1] =
+                      vim.treesitter.get_node_text(use_node, source)
+                  end
+                end
+              end
+            end
+          end
+          ret.env_override["USE_LIST"] = use_list
+        end
       end
 
       vim.api.nvim_win_set_cursor(0, { row, col })
@@ -102,6 +189,20 @@ return {
       end
 
       if in_test_cfg and env.MOD_ITEM_NAME == "tests" then
+        local test_fn_attrs = {
+          t("#[test]"),
+          t("#[tokio::test]"),
+        }
+
+        if Config.get("snippet.rust.rstest_support") == true then
+          local use_list = env["USE_LIST"] or {}
+          if vim.list_contains(use_list, "rstest::rstest") then
+            test_fn_attrs[#test_fn_attrs + 1] = t("#[rstest]")
+          elseif vim.list_contains(use_list, "rstest::*") then
+            test_fn_attrs[#test_fn_attrs + 1] = t("#[rstest]")
+          end
+        end
+
         -- function item
         return sn(
           nil,
@@ -121,10 +222,7 @@ return {
                 end
               end, { 2 }),
               name = i(1, "new_fn", { desc = "function name" }),
-              attr = c(2, {
-                t("#[test]"),
-                t("#[tokio::test]"),
-              }, { desc = "function attributes" }),
+              attr = c(2, test_fn_attrs, { desc = "function attributes" }),
               body = i(0),
             }
           )