refactor: injector function

Author: kawre

README.md

@@ -200,6 +200,22 @@ Whether to log [leetcode.nvim] status notifications
 logging = true
 ```
 
+### injector
+
+Inject code before or after your solution, injected code won't be submitted or run.
+
+```lua
+injector = { ---@type table<lc.lang, lc.inject>
+    ["cpp"] = {
+        before = { "#include <bits/stdc++.h>", "using namespace std;" },
+        after = "int main() {}",
+    },
+    ["java"] = {
+        before = "import java.util.*;",
+    },
+}
+```
+
 ### hooks
 
 List of functions that get executed on specified event
@@ -209,27 +225,11 @@ hooks = {
     ---@type fun()[]
     LeetEnter = {},
 
-    ---@type fun(question: { lang: string })[]
+    ---@type fun(question: lc.ui.Question)[]
     LeetQuestionNew = {},
 },
 ```
 
-### injector
-
-Inject code before or after your soluation, injected code won't be submitted or run.
-
-```lua
-injector = {
-    ["cpp"] = {
-        before = { "#include <bits/stdc++.h>", "using namespace std;" }, ---@type string|string[]
-        after = "int main(){}" ---@type string|string[]
-    },
-    ["java"] = {
-        before = "import java.util.*;", ---@type string|string[]
-    }
-}
-```
-
 ### image support
 
 Whether to render question description images using [image.nvim]

README.zh.md

@@ -99,6 +99,8 @@ https://github.com/kawre/leetcode.nvim/assets/69250723/aee6584c-e099-4409-b114-1
     ---@type boolean
     logging = true,
 
+    injector = {} ---@type table<lc.lang, lc.inject>
+
     cache = {
         update_interval = 60 * 60 * 24 * 7, ---@type integer 7 days
     },
@@ -152,8 +154,6 @@ https://github.com/kawre/leetcode.nvim/assets/69250723/aee6584c-e099-4409-b114-1
         focus_testcases = "H", ---@type string
         focus_result = "L", ---@type string
     },
-
-    injector = {} ---@type table<lc.lang, lc.inject>
 }
 ```
 
@@ -207,6 +207,22 @@ directory = vim.fn.stdpath("data") .. "/leetcode/"
 logging = true
 ```
 
+### injector
+
+在你的答案前后注入额外代码，注入的代码不会被提交或测试。
+
+```lua
+injector = { ---@type table<lc.lang, lc.inject>
+    ["cpp"] = {
+        before = { "#include <bits/stdc++.h>", "using namespace std;" },
+        after = "int main() {}",
+    },
+    ["java"] = {
+        before = "import java.util.*;",
+    },
+}
+```
+
 ### hooks
 
 在指定事件上执行的函数列表
@@ -216,27 +232,11 @@ hooks = {
     ---@type fun()[]
     LeetEnter = {},
 
-    ---@type fun(question: { lang: string })[]
+    ---@type fun(question: lc.ui.Question)[]
     LeetQuestionNew = {},
 },
 ```
 
-### injector
-
-在你的答案前后注入额外代码，注入的代码不会被提交或测试。
-
-```lua
-injector = {
-    ["cpp"] = {
-        before = { "#include <bits/stdc++.h>", "using namespace std;" }, ---@type string|string[]
-        after = "int main(){}" ---@type string|string[]
-    },
-    ["java"] = {
-        before = "import java.util.*;", ---@type string|string[]
-    }
-}
-```
-
 ### image support
 
 是否使用 [image.nvim] 渲染问题描述中的图片

lua/leetcode-ui/question.lua

@@ -25,12 +25,9 @@ function Question:get_snippet()
 
     local lang = utils.get_lang(self.lang)
     snip.code = (snip.code or ""):gsub("\r\n", "\n")
-    return ("%s%s @leet start\n%s\n%s @leet end%s"):format(
-        self:generate_inject_string(true),
-        lang.comment,
-        snip.code,
-        lang.comment,
-        self:generate_inject_string(false)
+
+    return self:injector(
+        ("%s @leet start\n%s\n%s @leet end"):format(lang.comment, snip.code, lang.comment)
     )
 end
 
@@ -44,16 +41,34 @@ function Question:create_file()
 end
 
 ---@private
----@param before boolean
-function Question:generate_inject_string(before)
+---@param code string
+function Question:injector(code)
     local injector = config.user.injector
-    if injector == nil or injector[self.lang] == nil then return "" end
-    local to_inject = before and injector[self.lang].before or injector[self.lang].after
-    if to_inject == nil or #to_inject == 0 then return "" end
 
-    if type(to_inject) ~= "string" then to_inject = table.concat(to_inject, "\n") end
+    local inject = injector[self.lang]
+    if not inject or vim.tbl_isempty(inject) then return code end
+
+    ---@param inj? string|string[]
+    ---@param before boolean
+    local function norm_inject(inj, before)
+        local res
+
+        if type(inj) == "table" then
+            res = table.concat(inj, "\n")
+        elseif type(inj) == "string" then
+            res = inj
+        end
+
+        if res and res ~= "" then
+            return before and (res .. "\n\n") or ("\n\n" .. res)
+        else
+            return ""
+        end
+    end
 
-    return before and to_inject .. "\n" or "\n" .. to_inject
+    return norm_inject(inject.before, true) --
+        .. code
+        .. norm_inject(inject.after, false)
 end
 
 Question.unmount = vim.schedule_wrap(function(self)

lua/leetcode/config/template.lua

@@ -59,6 +59,8 @@ local M = {
     ---@type boolean
     logging = true,
 
+    injector = {}, ---@type table<lc.lang, lc.inject>
+
     cache = {
         update_interval = 60 * 60 * 24 * 7, ---@type integer 7 days
     },
@@ -100,9 +102,6 @@ local M = {
         LeetQuestionNew = {},
     },
 
-    ---@type boolean
-    image_support = false, -- setting this to `true` will disable question description wrap
-
     keys = {
         toggle = { "q", "<Esc>" }, ---@type string|string[]
         confirm = { "<CR>" }, ---@type string|string[]
@@ -113,7 +112,8 @@ local M = {
         focus_result = "L", ---@type string
     },
 
-    injector = {} ---@type table<lc.lang, lc.inject>
+    ---@type boolean
+    image_support = false, -- setting this to `true` will disable question description wrap
 }
 
 return M