feat: Leet inject

README.md

@@ -314,6 +314,8 @@ image_support = false,
 
 - `restore` try to restore default question layout
 
+- `inject` re-inject code for the current question
+
 - `session`
 
  - `create` create a new session

lua/leetcode-ui/question.lua

@@ -51,35 +51,38 @@ function Question:create_file()
  return self.file:absolute()
 end
 
----@param code string
-function Question:injector(code)
- local injector = config.user.injector
- local inject = injector[self.lang] or {}
-
- ---@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
+---@param before boolean
+function Question:inject(before)
+ local inject = config.user.injector[self.lang] or {}
+ local inj = before and inject.before or inject.after
 
- if res and res ~= "" then
- return before and (res .. "\n\n") or ("\n\n" .. res)
- else
- return ""
- end
+ 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 nil
+ end
+end
+
+---@param code string
+function Question:injector(code)
  local lang = utils.get_lang(self.lang)
- return norm_inject(inject.before, true) --
+
+ local inj_before = self:inject(true) or ""
+ local inj_after = self:inject(false) or ""
+
+ return inj_before --
  .. ("%s @leet start\n"):format(lang.comment)
  .. code
  .. ("\n%s @leet end"):format(lang.comment)
- .. norm_inject(inject.after, false)
+ .. inj_after
 end
 
 ---@param pre? boolean
@@ -156,16 +159,17 @@ function Question:mount()
  return self
 end
 
+---@param inclusive? boolean
 ---@return integer, integer, string[]
-function Question:range()
+function Question:range(inclusive)
  local lines = vim.api.nvim_buf_get_lines(self.bufnr, 0, -1, false)
  local start_i, end_i = 1, #lines
 
  for i, line in ipairs(lines) do
  if line:match("@leet start") then
- start_i = i + 1
+ start_i = i + (inclusive and 0 or 1)
  elseif line:match("@leet end") then
- end_i = i - 1
+ end_i = i - (inclusive and 0 or 1)
  end
  end
 

lua/leetcode/command/init.lua

@@ -335,6 +335,27 @@ function cmd.restore()
  end
 end
 
+function cmd.inject()
+ local utils = require("leetcode.utils")
+ utils.auth_guard()
+ local q = utils.curr_question()
+ if not q then return end
+
+ local start_i, end_i = q:range(true)
+
+ if vim.api.nvim_buf_is_valid(q.bufnr) then
+ local before = q:inject(true)
+ if before then
+ vim.api.nvim_buf_set_lines(q.bufnr, 0, start_i - 1, false, vim.split(before, "\n"))
+ end
+
+ local after = q:inject(false)
+ if after then
+ vim.api.nvim_buf_set_lines(q.bufnr, end_i + 1, -1, false, vim.split(after, "\n"))
+ end
+ end
+end
+
 function cmd.get_active_session()
  local sessions = config.sessions.all
  return vim.tbl_filter(function(s) return s.is_active end, sessions)[1]
@@ -506,6 +527,7 @@ cmd.commands = {
  reset = { cmd.reset },
  last_submit = { cmd.last_submit },
  restore = { cmd.restore },
+ inject = { cmd.inject },
 
  session = {
  change = {

lua/leetcode/config/imports.lua

@@ -0,0 +1,73 @@
+---@class lc.Imports
+local imports = {}
+
+imports["python3"] = {
+ "from string import *",
+ "from re import *",
+ "from datetime import *",
+ "from collections import *",
+ "from heapq import *",
+ "from bisect import *",
+ "from copy import *",
+ "from math import *",
+ "from random import *",
+ "from statistics import *",
+ "from itertools import *",
+ "from functools import *",
+ "from operator import *",
+ "from io import *",
+ "from sys import *",
+ "from json import *",
+ "from builtins import *",
+ "import string",
+ "import re",
+ "import datetime",
+ "import collections",
+ "import heapq",
+ "import bisect",
+ "import copy",
+ "import math",
+ "import random",
+ "import statistics",
+ "import itertools",
+ "import functools",
+ "import operator",
+ "import io",
+ "import sys",
+ "import json",
+ "from typing import *",
+}
+
+imports["python"] = {
+ "from bisect import *",
+ "from collections import *",
+ "from copy import *",
+ "from datetime import *",
+ "from heapq import *",
+ "from math import *",
+ "from re import *",
+ "from string import *",
+ "from random import *",
+ "from itertools import *",
+ "from functools import *",
+ "from operator import *",
+ "from __builtin__ import *",
+ "import string",
+ "import re",
+ "import datetime",
+ "import collections",
+ "import heapq",
+ "import bisect",
+ "import copy",
+ "import math",
+ "import random",
+ "import itertools",
+ "import functools",
+ "import operator",
+}
+
+imports["java"] = {
+ "import java.util.*;",
+}
+
+return imports

lua/leetcode/config/init.lua

@@ -28,6 +28,7 @@ local config = {
  icons = require("leetcode.config.icons"),
  sessions = require("leetcode.config.sessions"),
  stats = require("leetcode.config.stats"),
+ imports = require("leetcode.config.imports"),
 
  ---@type lc.UserStatus
  auth = {}, ---@diagnostic disable-line