some refactors

This commit is contained in:
Barrett Ruth 2025-10-03 14:34:49 -04:00
parent fea9835436
commit 1520939d4b
5 changed files with 249 additions and 163 deletions

View file

@ -79,29 +79,22 @@ end
---@param platform string
---@param contest_id string
---@return ContestData?
---@return ContestData
function M.get_contest_data(platform, contest_id)
vim.validate({
platform = { platform, 'string' },
contest_id = { contest_id, 'string' },
})
if not cache_data[platform] then
return nil
end
local contest_data = cache_data[platform][contest_id]
if not contest_data or vim.tbl_isempty(contest_data) then
return nil
end
return contest_data
return cache_data[platform][contest_id] or {}
end
---@param platform string
---@param contest_id string
---@param problems Problem[]
function M.set_contest_data(platform, contest_id, problems)
---@param contest_name? string
---@param display_name? string
function M.set_contest_data(platform, contest_id, problems, contest_name, display_name)
vim.validate({
platform = { platform, 'string' },
contest_id = { contest_id, 'string' },
@ -109,36 +102,17 @@ function M.set_contest_data(platform, contest_id, problems)
})
cache_data[platform] = cache_data[platform] or {}
local existing = cache_data[platform][contest_id] or {}
local existing_by_id = {}
if existing.problems then
for _, p in ipairs(existing.problems) do
existing_by_id[p.id] = p
end
local out = {
name = contest_name,
display_name = display_name,
problems = vim.deepcopy(problems),
index_map = {},
}
for i, p in ipairs(out.problems) do
out.index_map[p.id] = i
end
local merged = {}
for _, p in ipairs(problems) do
local prev = existing_by_id[p.id] or {}
local merged_p = {
id = p.id,
name = p.name or prev.name,
test_cases = prev.test_cases,
timeout_ms = prev.timeout_ms,
memory_mb = prev.memory_mb,
interactive = prev.interactive,
}
table.insert(merged, merged_p)
end
existing.problems = merged
existing.index_map = {}
for i, p in ipairs(merged) do
existing.index_map[p.id] = i
end
cache_data[platform][contest_id] = existing
cache_data[platform][contest_id] = out
M.save()
end

View file

@ -1,67 +1,109 @@
local M = {}
local utils = require('cp.utils')
local logger = require('cp.log')
local utils = require('cp.utils')
local function syshandle(result)
if result.code ~= 0 then
local msg = 'Scraper failed: ' .. (result.stderr or 'Unknown error')
logger.log(msg, vim.log.levels.ERROR)
return {
success = false,
error = msg,
}
return { success = false, error = msg }
end
local ok, data = pcall(vim.json.decode, result.stdout)
if not ok then
local msg = 'Failed to parse scraper output: ' .. tostring(data)
logger.log(msg, vim.log.levels.ERROR)
return {
success = false,
error = msg,
}
return { success = false, error = msg }
end
return {
success = true,
data = data,
}
return { success = true, data = data }
end
---@param platform string
---@param subcommand string
---@param args string[]
---@param opts { sync?: boolean, ndjson?: boolean, on_event?: fun(ev: table), on_exit?: fun(result: table) }
local function run_scraper(platform, subcommand, args, opts)
if not utils.setup_python_env() then
local msg = 'Python environment setup failed'
logger.log(msg, vim.log.levels.ERROR)
return {
success = false,
message = msg,
}
end
local plugin_path = utils.get_plugin_path()
local cmd = {
'uv',
'run',
'--directory',
plugin_path,
'-m',
'scrapers.' .. platform,
subcommand,
}
local cmd = { 'uv', 'run', '--directory', plugin_path, '-m', 'scrapers.' .. platform, subcommand }
vim.list_extend(cmd, args)
local sysopts = {
text = true,
timeout = 30000,
}
if opts and opts.ndjson then
local uv = vim.loop
local stdout = uv.new_pipe(false)
local stderr = uv.new_pipe(false)
local buf = ''
if opts.sync then
local handle = uv.spawn(
cmd[1],
{ args = vim.list_slice(cmd, 2), stdio = { nil, stdout, stderr } },
function(code, signal)
if buf ~= '' and opts.on_event then
local ok_tail, ev_tail = pcall(vim.json.decode, buf)
if ok_tail then
opts.on_event(ev_tail)
end
buf = ''
end
if opts.on_exit then
opts.on_exit({ success = (code == 0), code = code, signal = signal })
end
if not stdout:is_closing() then
stdout:close()
end
if not stderr:is_closing() then
stderr:close()
end
if handle and not handle:is_closing() then
handle:close()
end
end
)
if not handle then
logger.log('Failed to start scraper process', vim.log.levels.ERROR)
return { success = false, error = 'spawn failed' }
end
uv.read_start(stdout, function(_, data)
if data == nil then
if buf ~= '' and opts.on_event then
local ok_tail, ev_tail = pcall(vim.json.decode, buf)
if ok_tail then
opts.on_event(ev_tail)
end
buf = ''
end
return
end
buf = buf .. data
while true do
local s, e = buf:find('\n', 1, true)
if not s then
break
end
local line = buf:sub(1, s - 1)
buf = buf:sub(e + 1)
local ok, ev = pcall(vim.json.decode, line)
if ok and opts.on_event then
opts.on_event(ev)
end
end
end)
uv.read_start(stderr, function(_, _) end)
return
end
local sysopts = { text = true, timeout = 30000 }
if opts and opts.sync then
local result = vim.system(cmd, sysopts):wait()
return syshandle(result)
else
vim.system(cmd, sysopts, function(result)
return opts.on_exit(syshandle(result))
if opts and opts.on_exit then
return opts.on_exit(syshandle(result))
end
end)
end
end
@ -93,41 +135,48 @@ end
function M.scrape_contest_list(platform)
local result = run_scraper(platform, 'contests', {}, { sync = true })
if not result.success or not result.data.contests then
if not result or not result.success or not (result.data and result.data.contests) then
logger.log(
('Could not scrape contests list for platform %s: %s'):format(platform, result.msg),
('Could not scrape contests list for platform %s: %s'):format(
platform,
(result and result.error) or 'unknown'
),
vim.log.levels.ERROR
)
return {}
end
return result.data.contests
end
function M.scrape_problem_tests(platform, contest_id, problem_id, callback)
run_scraper(platform, 'tests', { contest_id, problem_id }, {
on_exit = function(result)
if not result.success or not result.data.tests then
logger.log(
'Failed to load tests: ' .. (result.msg or 'unknown error'),
vim.log.levels.ERROR
)
return {}
---@param platform string
---@param contest_id string
---@param callback fun(data: table)|nil
function M.scrape_all_tests(platform, contest_id, callback)
run_scraper(platform, 'tests', { contest_id }, {
ndjson = true,
on_event = function(ev)
if ev.done then
return
end
if ev.error and ev.problem_id then
logger.log(
('Failed to load tests for %s/%s: %s'):format(contest_id, ev.problem_id, ev.error),
vim.log.levels.WARN
)
return
end
if not ev.problem_id or not ev.tests then
return
end
vim.schedule(function()
vim.system({ 'mkdir', '-p', 'build', 'io' }):wait()
local config = require('cp.config')
local base_name = config.default_filename(contest_id, problem_id)
for i, test_case in ipairs(result.data.tests) do
local base_name = config.default_filename(contest_id, ev.problem_id)
for i, t in ipairs(ev.tests) do
local input_file = 'io/' .. base_name .. '.' .. i .. '.cpin'
local expected_file = 'io/' .. base_name .. '.' .. i .. '.cpout'
local input_content = test_case.input:gsub('\r', '')
local expected_content = test_case.expected:gsub('\r', '')
local input_content = t.input:gsub('\r', '')
local expected_content = t.expected:gsub('\r', '')
pcall(vim.fn.writefile, vim.split(input_content, '\n', { trimempty = true }), input_file)
pcall(
vim.fn.writefile,
@ -136,7 +185,13 @@ function M.scrape_problem_tests(platform, contest_id, problem_id, callback)
)
end
if type(callback) == 'function' then
callback(result.data)
callback({
tests = ev.tests,
timeout_ms = ev.timeout_ms or 0,
memory_mb = ev.memory_mb or 0,
interactive = ev.interactive or false,
problem_id = ev.problem_id,
})
end
end)
end,

View file

@ -28,45 +28,26 @@ function M.set_platform(platform)
return true
end
local function backfill_missing_tests(platform, contest_id, problems)
cache.load()
local missing = {}
for _, prob in ipairs(problems) do
if not cache.get_test_cases(platform, contest_id, prob.id) then
table.insert(missing, prob.id)
end
end
if #missing == 0 then
logger.log(('All problems already cached for %s contest %s.'):format(platform, contest_id))
return
end
for _, pid in ipairs(missing) do
local captured = pid
scraper.scrape_problem_tests(platform, contest_id, captured, function(result)
local cached_tests = {}
if result.tests then
for i, t in ipairs(result.tests) do
cached_tests[i] = { index = i, input = t.input, expected = t.expected }
end
end
cache.set_test_cases(
platform,
contest_id,
captured,
cached_tests,
result.timeout_ms,
result.memory_mb
)
end)
end
end
---@class TestCaseLite
---@field input string
---@field expected string
---@class ScrapeEvent
---@field problem_id string
---@field tests TestCaseLite[]|nil
---@field timeout_ms integer|nil
---@field memory_mb integer|nil
---@field interactive boolean|nil
---@field error string|nil
---@field done boolean|nil
---@field succeeded integer|nil
---@field failed integer|nil
---@param platform string
---@param contest_id string
---@param language string|nil
---@param problem_id string|nil
function M.setup_contest(platform, contest_id, language, problem_id)
if not platform then
logger.log('No platform configured. Use :CP <platform> <contest> [--{lang=<lang>,debug} first.')
return
end
local config = config_module.get_config()
if not vim.tbl_contains(config.scrapers, platform) then
logger.log(('Scraping disabled for %s.'):format(platform), vim.log.levels.WARN)
@ -75,27 +56,70 @@ function M.setup_contest(platform, contest_id, language, problem_id)
state.set_contest_id(contest_id)
cache.load()
local contest_data = cache.get_contest_data(platform, contest_id)
local contest_data = cache.get_contest_data(platform, contest_id)
if not contest_data or not contest_data.problems then
logger.log('Fetching contests problems...', vim.log.levels.INFO, true)
scraper.scrape_contest_metadata(platform, contest_id, function(result)
local problems = result.problems or {}
cache.set_contest_data(platform, contest_id, problems)
cache.set_contest_data(platform, contest_id, problems, result.name, result.display_name)
logger.log(('Found %d problems for %s contest %s.'):format(#problems, platform, contest_id))
local pid = problem_id or (problems[1] and problems[1].id)
if pid then
M.setup_problem(pid, language)
end
backfill_missing_tests(platform, contest_id, problems)
end)
else
local problems = contest_data.problems
local pid = problem_id or (problems[1] and problems[1].id)
if pid then
contest_data = cache.get_contest_data(platform, contest_id)
local pid = contest_data.problems[problem_id and contest_data.index_map[problem_id] or 1].id
M.setup_problem(pid, language)
end
backfill_missing_tests(platform, contest_id, problems)
local cached_len = #vim.tbl_filter(function(p)
return cache.get_test_cases(platform, contest_id, p.id) ~= nil
end, problems)
if cached_len < #problems then
scraper.scrape_all_tests(platform, contest_id, function(ev)
if not ev or not ev.tests or not ev.problem_id then
return
end
local cached_tests = {}
for i, t in ipairs(ev.tests) do
cached_tests[i] = { index = i, input = t.input, expected = t.expected }
end
cache.set_test_cases(
platform,
contest_id,
ev.problem_id,
cached_tests,
ev.timeout_ms or 0,
ev.memory_mb or 0
)
end)
end
end)
return
end
local problems = contest_data.problems
local pid = problems[(problem_id and contest_data.index_map[problem_id] or 1)].id
M.setup_problem(pid, language)
local cached_len = #vim.tbl_filter(function(p)
return cache.get_test_cases(platform, contest_id, p.id) ~= nil
end, problems)
if cached_len < #problems then
scraper.scrape_all_tests(platform, contest_id, function(ev)
if not ev or not ev.tests or not ev.problem_id then
return
end
local cached_tests = {}
for i, t in ipairs(ev.tests) do
cached_tests[i] = { index = i, input = t.input, expected = t.expected }
end
cache.set_test_cases(
platform,
contest_id,
ev.problem_id,
cached_tests,
ev.timeout_ms or 0,
ev.memory_mb or 0
)
end)
end
end
@ -195,19 +219,9 @@ function M.navigate_problem(direction, language)
end
local problems = contest_data.problems
local current_index
for i, prob in ipairs(problems) do
if prob.id == current_problem_id then
current_index = i
break
end
end
if not current_index then
M.setup_contest(platform, contest_id, language, problems[1].id)
return
end
local index = contest_data.index_map[current_problem_id]
local new_index = current_index + direction
local new_index = index + direction
if new_index < 1 or new_index > #problems then
return
end