diff --git a/doc/cp.txt b/doc/cp.txt index 7966559..ca856c6 100644 --- a/doc/cp.txt +++ b/doc/cp.txt @@ -22,17 +22,28 @@ Optional: COMMANDS *cp-commands* *:CP* -:CP {contest} Set up contest environment for {contest}. - Available contests: atcoder, codeforces, cses +cp.nvim uses a single :CP command with intelligent argument parsing: -:CP {contest} {contest_id} {problem_letter?} - Set up problem from {contest}. Extract contest_id - and problem_letter (optional) from the problem URL. - Scrapes test cases and creates source file. +Setup Commands ~ -:CP {contest_id} {problem_letter} - Set up problem in current contest mode. - Requires contest to be set first with :CP {contest} +:CP {platform} {contest_id} {problem_id} + Full setup: set platform, load contest metadata, + and set up specific problem. Scrapes test cases + and creates source file. + Example: :CP codeforces 1933 a + +:CP {platform} {contest_id} Contest setup: set platform and load contest + metadata for navigation. Caches problem list. + Example: :CP atcoder abc324 + +:CP {platform} Platform setup: set platform only. + Example: :CP cses + +:CP {problem_id} Problem switch: switch to different problem + within current contest context. + Example: :CP b (switch to problem b) + +Action Commands ~ :CP run Compile and run current problem with test input. Shows execution time and output comparison. @@ -43,64 +54,43 @@ COMMANDS *cp-commands* :CP diff Enter diff mode to compare actual vs expected output. Run again to exit diff mode. +Navigation Commands ~ + +:CP next Navigate to next problem in current contest. + Stops at last problem (no wrapping). + +:CP prev Navigate to previous problem in current contest. + Stops at first problem (no wrapping). + CONFIGURATION *cp-config* -cp.nvim is automatically lazy-loaded - no config/setup is required. +cp.nvim works out of the box. No setup required. -Provide extra options via a setup() function with your package manager. For -example, with lazy.nvim (https://github.com/folke/lazy.nvim): - - { - 'barrett-ruth/cp.nvim', - config = function() - local ls = require('luasnip') - local s = ls.snippet - - require('cp').setup({ - debug = false, - contests = { - default = { - cpp_version = 20, - compile_flags = { "-O2", "-DLOCAL", "-Wall", "-Wextra" }, - debug_flags = { "-g3", "-fsanitize=address,undefined", "-DLOCAL" }, - timeout_ms = 2000, - }, - atcoder = { - cpp_version = 23, - }, +Optional: > + vim.g.cp = { + config = { + debug = false, + contests = { + default = { + cpp_version = 20, + compile_flags = { "-O2", "-DLOCAL", "-Wall", "-Wextra" }, + debug_flags = { "-g3", "-fsanitize=address,undefined", "-DLOCAL" }, + timeout_ms = 2000, }, - snippets = { - cses = { - s("cses", "#include \nusing namespace std;\n\nint main() {\n\t$0\n}") - }, - }, - hooks = { - before_run = function(problem_id) - vim.cmd.w() - vim.lsp.buf.format() - end, - before_debug = function(problem_id) - ... - end - }, - tile = function(source_buf, input_buf, output_buf) - vim.api.nvim_set_current_buf(source_buf) - vim.cmd.vsplit() - vim.api.nvim_set_current_buf(output_buf) - vim.cmd.vsplit() - vim.api.nvim_set_current_buf(input_buf) - vim.cmd('wincmd h | wincmd h') - end, - filename = function(contest, problem_id, problem_letter) - if contest == "atcoder" then - return problem_id:lower() .. (problem_letter or "") .. ".cpp" - else - return problem_id:lower() .. (problem_letter or "") .. ".cc" - end - end, - }) - end + atcoder = { cpp_version = 23 }, + }, + hooks = { + before_run = function(problem_id) vim.cmd.w() end, + }, + tile = function(source_buf, input_buf, output_buf) + -- custom window layout + end, + filename = function(contest, problem_id, problem_letter) + -- custom filename generation + end, + } } +< Configuration options: @@ -145,70 +135,96 @@ AtCoder ~ URL format: https://atcoder.jp/contests/abc123/tasks/abc123_a In terms of cp.nvim, this corresponds to: - +- Platform: atcoder - Contest ID: abc123 -- Problem letter: a +- Problem ID: a Usage examples: > - :CP atcoder abc123 a " Set up problem A from contest ABC123 - :CP atcoder " Set up AtCoder contest mode first - :CP abc123 a " Then set up problem (if contest mode is set) + :CP atcoder abc123 a " Full setup: problem A from contest ABC123 + :CP atcoder abc123 " Contest setup: load contest metadata only + :CP b " Switch to problem B (if contest loaded) + :CP next " Navigate to next problem in contest < Codeforces ~ *cp-codeforces* URL format: https://codeforces.com/contest/1234/problem/A In terms of cp.nvim, this corresponds to: - +- Platform: codeforces - Contest ID: 1234 -- Problem letter: A +- Problem ID: a (lowercase) Usage examples: > - :CP codeforces 1234 a " Set up problem A from contest id 1234 - :CP codeforces " Set up Codeforces contest mode first - :CP 1234 a " Then set up problem (if contest mode is set) + :CP codeforces 1934 a " Full setup: problem A from contest 1934 + :CP codeforces 1934 " Contest setup: load contest metadata only + :CP c " Switch to problem C (if contest loaded) + :CP prev " Navigate to previous problem in contest < CSES ~ *cp-cses* URL format: https://cses.fi/problemset/task/1068 +CSES is organized by categories rather than contests. Currently all problems +are grouped under "CSES Problem Set" category. + In terms of cp.nvim, this corresponds to: -- Problem ID: 1068 +- Platform: cses +- Contest ID: "CSES Problem Set" (category) +- Problem ID: 1068 (numeric) Usage examples: > - :CP cses 1068 " Set up problem 1068 - :CP cses " Set up CSES contest mode first - :CP 1068 " Then set up problem (if contest mode is set) + :CP cses 1068 " Set up problem 1068 from CSES + :CP 1070 " Switch to problem 1070 (if CSES loaded) + :CP next " Navigate to next problem in CSES < COMPLETE WORKFLOW EXAMPLE *cp-example* -Example: Setting up AtCoder problem ABC123-A +Example: Setting up and solving AtCoder contest ABC324 -1. Browse to https://atcoder.jp/contests/abc123/tasks/abc123_a -2. Read the problem statement on the website -3. In Neovim, extract identifiers and set up - :CP atcoder abc123 a -< - This creates abc123a.cc (or however you've configured the filename in - *cp-setup*) and scrapes test cases -4. Code and test +1. Browse to https://atcoder.jp/contests/abc324 +2. Set up contest and load metadata: > + :CP atcoder abc324 +< This caches all problems (A, B, C, D, E, F, G) for navigation + +3. Start with problem A: > + :CP a +< This creates abc324a.cc and scrapes test cases + +4. Code your solution, then test: > :CP run -5. Debug +< +5. If needed, debug: > :CP debug -6. Test: > +< +6. Compare output visually: > :CP diff -7. Submit remote on AtCoder +< +7. Move to next problem: > + :CP next +< This automatically sets up problem B + +8. Continue solving problems with :CP next/:CP prev navigation +9. Submit solutions on AtCoder website + +Example: Quick setup for single Codeforces problem > + :CP codeforces 1933 a " One command setup + :CP run " Test immediately +< FILE STRUCTURE *cp-files* -cp.nvim creates the following file structure upon setup: +cp.nvim creates the following file structure upon problem setup: - problem.cc - build/*.{run,debug} + {contest_id}{problem_id}.cc " Source file (e.g. abc324a.cc) + build/ + {contest_id}{problem_id}.run " Compiled binary io/ - problem.in - problem.out - problem.expected + {contest_id}{problem_id}.in " Test input + {contest_id}{problem_id}.out " Program output + {contest_id}{problem_id}.expected " Expected output + +The plugin automatically manages this structure and navigation between problems +maintains proper file associations. SNIPPETS *cp-snippets* diff --git a/lua/cp/cache.lua b/lua/cp/cache.lua new file mode 100644 index 0000000..e74ace0 --- /dev/null +++ b/lua/cp/cache.lua @@ -0,0 +1,90 @@ +local M = {} + +local cache_file = vim.fn.stdpath("data") .. "/cp-contest-cache.json" +local cache_data = {} + +local function get_expiry_date(platform) + if platform == "cses" then + return os.time() + (30 * 24 * 60 * 60) + end + return nil +end + +local function is_cache_valid(contest_data, platform) + if platform ~= "cses" then + return true + end + + local expires_at = contest_data.expires_at + if not expires_at then + return false + end + + return os.time() < expires_at +end + +function M.load() + if vim.fn.filereadable(cache_file) == 0 then + cache_data = {} + return + end + + local content = vim.fn.readfile(cache_file) + if #content == 0 then + cache_data = {} + return + end + + local ok, decoded = pcall(vim.json.decode, table.concat(content, "\n")) + if ok then + cache_data = decoded + else + cache_data = {} + end +end + +function M.save() + vim.fn.mkdir(vim.fn.fnamemodify(cache_file, ":h"), "p") + local encoded = vim.json.encode(cache_data) + vim.fn.writefile(vim.split(encoded, "\n"), cache_file) +end + +function M.get_contest_data(platform, contest_id) + if not cache_data[platform] then + return nil + end + + local contest_data = cache_data[platform][contest_id] + if not contest_data then + return nil + end + + if not is_cache_valid(contest_data, platform) then + return nil + end + + return contest_data +end + +function M.set_contest_data(platform, contest_id, problems) + if not cache_data[platform] then + cache_data[platform] = {} + end + + cache_data[platform][contest_id] = { + problems = problems, + scraped_at = os.date("%Y-%m-%d"), + expires_at = get_expiry_date(platform), + } + + M.save() +end + +function M.clear_contest_data(platform, contest_id) + if cache_data[platform] and cache_data[platform][contest_id] then + cache_data[platform][contest_id] = nil + M.save() + end +end + +return M diff --git a/lua/cp/config.lua b/lua/cp/config.lua index 548bdc9..9f39c60 100644 --- a/lua/cp/config.lua +++ b/lua/cp/config.lua @@ -44,8 +44,8 @@ local function extend_contest_config(base_config, contest_config) local result = vim.tbl_deep_extend("force", base_config, contest_config) local std_flag = ("-std=c++%d"):format(result.cpp_version) - table.insert(result.compile_flags, 1, std_flag) - table.insert(result.debug_flags, 1, std_flag) + result.compile_flags = vim.list_extend({ std_flag }, result.compile_flags) + result.debug_flags = vim.list_extend({ std_flag }, result.debug_flags) return result end diff --git a/lua/cp/execute.lua b/lua/cp/execute.lua index d1acc8e..774346b 100644 --- a/lua/cp/execute.lua +++ b/lua/cp/execute.lua @@ -2,12 +2,20 @@ local M = {} local signal_codes = { [128] = "SIGILL", - [130] = "SIGABRT", - [131] = "SIGBUS", + [130] = "SIGINT", + [131] = "SIGQUIT", + [132] = "SIGILL", + [133] = "SIGTRAP", + [134] = "SIGABRT", + [135] = "SIGBUS", [136] = "SIGFPE", - [135] = "SIGSEGV", - [137] = "SIGPIPE", - [139] = "SIGTERM", + [137] = "SIGKILL", + [138] = "SIGUSR1", + [139] = "SIGSEGV", + [140] = "SIGUSR2", + [141] = "SIGPIPE", + [142] = "SIGALRM", + [143] = "SIGTERM", } local function ensure_directories() @@ -31,34 +39,41 @@ local function execute_binary(binary_path, input_data, timeout_ms) local end_time = vim.loop.hrtime() local execution_time = (end_time - start_time) / 1000000 + local actual_code = result.code or 0 + return { stdout = result.stdout or "", stderr = result.stderr or "", - code = result.code, + code = actual_code, time_ms = execution_time, timed_out = result.code == 124, } end local function format_output(exec_result, expected_file, is_debug) - local lines = { exec_result.stdout } + local output_lines = { exec_result.stdout } + local metadata_lines = {} if exec_result.timed_out then - table.insert(lines, "\n[code]: 124 (TIMEOUT)") + table.insert(metadata_lines, "[code]: 124 (TIMEOUT)") elseif exec_result.code >= 128 then local signal_name = signal_codes[exec_result.code] or "SIGNAL" - table.insert(lines, ("\n[code]: %d (%s)"):format(exec_result.code, signal_name)) + table.insert(metadata_lines, ("[code]: %d (%s)"):format(exec_result.code, signal_name)) else - table.insert(lines, ("\n[code]: %d"):format(exec_result.code)) + table.insert(metadata_lines, ("[code]: %d"):format(exec_result.code)) end - table.insert(lines, ("\n[time]: %.2f ms"):format(exec_result.time_ms)) - table.insert(lines, ("\n[debug]: %s"):format(is_debug and "true" or "false")) + table.insert(metadata_lines, ("[time]: %.2f ms"):format(exec_result.time_ms)) + table.insert(metadata_lines, ("[debug]: %s"):format(is_debug and "true" or "false")) if vim.fn.filereadable(expected_file) == 1 and exec_result.code == 0 then local expected_content = vim.fn.readfile(expected_file) local actual_lines = vim.split(exec_result.stdout, "\n") + while #actual_lines > 0 and actual_lines[#actual_lines] == "" do + table.remove(actual_lines) + end + local matches = #actual_lines == #expected_content if matches then for i, line in ipairs(actual_lines) do @@ -69,10 +84,10 @@ local function format_output(exec_result, expected_file, is_debug) end end - table.insert(lines, ("\n[matches]: %s"):format(matches and "true" or "false")) + table.insert(metadata_lines, ("[matches]: %s"):format(matches and "true" or "false")) end - return table.concat(lines, "") + return table.concat(output_lines, "") .. "\n" .. table.concat(metadata_lines, "\n") end ---@param ctx ProblemContext diff --git a/lua/cp/health.lua b/lua/cp/health.lua index 5a833ff..d902934 100644 --- a/lua/cp/health.lua +++ b/lua/cp/health.lua @@ -60,17 +60,19 @@ local function check_luasnip() end local function check_config() - local cp = require("cp") - if cp.is_initialized() then - vim.health.ok("Plugin initialized") + vim.health.ok("Plugin ready") - if vim.g.cp_contest then - vim.health.info("Current contest: " .. vim.g.cp_contest) - else - vim.health.info("No contest mode set") + if vim.g.cp and vim.g.cp.platform then + local info = vim.g.cp.platform + if vim.g.cp.contest_id then + info = info .. " " .. vim.g.cp.contest_id + if vim.g.cp.problem_id then + info = info .. " " .. vim.g.cp.problem_id + end end + vim.health.info("Current context: " .. info) else - vim.health.warn("Plugin not initialized - configuration may be incomplete") + vim.health.info("No contest context set") end end diff --git a/lua/cp/init.lua b/lua/cp/init.lua index 5dc988b..8abfcf4 100644 --- a/lua/cp/init.lua +++ b/lua/cp/init.lua @@ -1,69 +1,87 @@ -local config_module = require("cp.config") -local snippets = require("cp.snippets") -local execute = require("cp.execute") -local scrape = require("cp.scrape") -local window = require("cp.window") -local logger = require("cp.log") -local problem = require("cp.problem") - local M = {} local config = {} -if not vim.fn.has("nvim-0.10.0") then - logger.log("cp.nvim requires nvim-0.10.0+", vim.log.levels.ERROR) - return M +local config_module, snippets, execute, scrape, window, logger, problem, cache + +local function lazy_require() + if not config_module then + if not vim.fn.has("nvim-0.10.0") then + vim.notify("[cp.nvim]: requires nvim-0.10.0+", vim.log.levels.ERROR) + return false + end + + config_module = require("cp.config") + snippets = require("cp.snippets") + execute = require("cp.execute") + scrape = require("cp.scrape") + window = require("cp.window") + logger = require("cp.log") + problem = require("cp.problem") + cache = require("cp.cache") + end + return true end -local competition_types = { "atcoder", "codeforces", "cses" } +local platforms = { "atcoder", "codeforces", "cses" } +local actions = { "run", "debug", "diff", "next", "prev" } -local function setup_contest(contest_type) - if not vim.tbl_contains(competition_types, contest_type) then - logger.log( - ("unknown contest type. Available: [%s]"):format(table.concat(competition_types, ", ")), - vim.log.levels.ERROR - ) - return +local function set_platform(platform) + if not vim.tbl_contains(platforms, platform) then + logger.log(("unknown platform. Available: [%s]"):format(table.concat(platforms, ", ")), vim.log.levels.ERROR) + return false end - vim.g.cp_contest = contest_type + vim.g.cp = vim.g.cp or {} + vim.g.cp.platform = platform vim.fn.mkdir("build", "p") vim.fn.mkdir("io", "p") - logger.log(("set up %s contest environment"):format(contest_type)) + return true end +---@param contest_id string +---@param problem_id? string local function setup_problem(contest_id, problem_id) - if not vim.g.cp_contest then - logger.log("no contest mode set. run :CP first", vim.log.levels.ERROR) + if not vim.g.cp or not vim.g.cp.platform then + logger.log("no platform set. run :CP first", vim.log.levels.ERROR) return end - if vim.g.cp_diff_mode then + local metadata_result = scrape.scrape_contest_metadata(vim.g.cp.platform, contest_id) + if not metadata_result.success then + logger.log( + "failed to load contest metadata: " .. (metadata_result.error or "unknown error"), + vim.log.levels.WARN + ) + end + + if vim.g.cp and vim.g.cp.diff_mode then vim.cmd.diffoff() - if vim.g.cp_saved_session then - vim.fn.delete(vim.g.cp_saved_session) - vim.g.cp_saved_session = nil + if vim.g.cp.saved_session then + vim.fn.delete(vim.g.cp.saved_session) + vim.g.cp.saved_session = nil end - if vim.g.cp_temp_output then - vim.fn.delete(vim.g.cp_temp_output) - vim.g.cp_temp_output = nil + if vim.g.cp.temp_output then + vim.fn.delete(vim.g.cp.temp_output) + vim.g.cp.temp_output = nil end - vim.g.cp_diff_mode = false + vim.g.cp.diff_mode = false end vim.cmd("silent only") - vim.g.cp_contest_id = contest_id - vim.g.cp_problem_id = problem_id + vim.g.cp.contest_id = contest_id + vim.g.cp.problem_id = problem_id - local ctx = problem.create_context(vim.g.cp_contest, contest_id, problem_id, config) + local ctx = problem.create_context(vim.g.cp.platform, contest_id, problem_id, config) local scrape_result = scrape.scrape_problem(ctx) if not scrape_result.success then - logger.log("scraping failed: " .. scrape_result.error, vim.log.levels.WARN) + logger.log("scraping failed: " .. (scrape_result.error or "unknown error"), vim.log.levels.WARN) logger.log("you can manually add test cases to io/ directory", vim.log.levels.INFO) else - logger.log(("scraped %d test case(s) for %s"):format(scrape_result.test_count, scrape_result.problem_id)) + local test_count = scrape_result.test_count or 0 + logger.log(("scraped %d test case(s) for %s"):format(test_count, scrape_result.problem_id)) end vim.cmd.e(ctx.source_file) @@ -71,8 +89,8 @@ local function setup_problem(contest_id, problem_id) if vim.api.nvim_buf_get_lines(0, 0, -1, true)[1] == "" then local has_luasnip, luasnip = pcall(require, "luasnip") if has_luasnip then - vim.api.nvim_buf_set_lines(0, 0, -1, false, { vim.g.cp_contest }) - vim.api.nvim_win_set_cursor(0, { 1, #vim.g.cp_contest }) + vim.api.nvim_buf_set_lines(0, 0, -1, false, { vim.g.cp.platform }) + vim.api.nvim_win_set_cursor(0, { 1, #vim.g.cp.platform }) vim.cmd.startinsert({ bang = true }) vim.schedule(function() @@ -82,7 +100,7 @@ local function setup_problem(contest_id, problem_id) vim.cmd.stopinsert() end) else - vim.api.nvim_input(("i%s"):format(vim.g.cp_contest)) + vim.api.nvim_input(("i%s"):format(vim.g.cp.platform)) end end @@ -123,15 +141,15 @@ local function run_problem() config.hooks.before_run(problem_id) end - if not vim.g.cp_contest then - logger.log("no contest mode set", vim.log.levels.ERROR) + if not vim.g.cp or not vim.g.cp.platform then + logger.log("no platform set", vim.log.levels.ERROR) return end - local contest_config = config.contests[vim.g.cp_contest] + local contest_config = config.contests[vim.g.cp.platform] vim.schedule(function() - local ctx = problem.create_context(vim.g.cp_contest, vim.g.cp_contest_id, vim.g.cp_problem_id, config) + local ctx = problem.create_context(vim.g.cp.platform, vim.g.cp.contest_id, vim.g.cp.problem_id, config) execute.run_problem(ctx, contest_config, false) vim.cmd.checktime() end) @@ -147,26 +165,26 @@ local function debug_problem() config.hooks.before_debug(problem_id) end - if not vim.g.cp_contest then - logger.log("no contest mode set", vim.log.levels.ERROR) + if not vim.g.cp or not vim.g.cp.platform then + logger.log("no platform set", vim.log.levels.ERROR) return end - local contest_config = config.contests[vim.g.cp_contest] + local contest_config = config.contests[vim.g.cp.platform] vim.schedule(function() - local ctx = problem.create_context(vim.g.cp_contest, vim.g.cp_contest_id, vim.g.cp_problem_id, config) + local ctx = problem.create_context(vim.g.cp.platform, vim.g.cp.contest_id, vim.g.cp.problem_id, config) execute.run_problem(ctx, contest_config, true) vim.cmd.checktime() end) end local function diff_problem() - if vim.g.cp_diff_mode then + if vim.g.cp and vim.g.cp.diff_mode then local tile_fn = config.tile or window.default_tile - window.restore_layout(vim.g.cp_saved_layout, tile_fn) - vim.g.cp_diff_mode = false - vim.g.cp_saved_layout = nil + window.restore_layout(vim.g.cp.saved_layout, tile_fn) + vim.g.cp.diff_mode = false + vim.g.cp.saved_layout = nil logger.log("exited diff mode") else local problem_id = get_current_problem() @@ -174,80 +192,196 @@ local function diff_problem() return end - local ctx = problem.create_context(vim.g.cp_contest, vim.g.cp_contest_id, vim.g.cp_problem_id, config) + local ctx = problem.create_context(vim.g.cp.platform, vim.g.cp.contest_id, vim.g.cp.problem_id, config) if vim.fn.filereadable(ctx.expected_file) == 0 then logger.log(("No expected output file found: %s"):format(ctx.expected_file), vim.log.levels.ERROR) return end - vim.g.cp_saved_layout = window.save_layout() + vim.g.cp = vim.g.cp or {} + vim.g.cp.saved_layout = window.save_layout() local result = vim.system({ "awk", "/^\\[[^]]*\\]:/ {exit} {print}", ctx.output_file }, { text = true }):wait() local actual_output = result.stdout window.setup_diff_layout(actual_output, ctx.expected_file, ctx.input_file) - vim.g.cp_diff_mode = true + vim.g.cp.diff_mode = true logger.log("entered diff mode") end end -local initialized = false - -function M.is_initialized() - return initialized -end - -function M.setup(user_config) - if initialized and not user_config then +---@param delta number 1 for next, -1 for prev +local function navigate_problem(delta) + if not vim.g.cp or not vim.g.cp.platform or not vim.g.cp.contest_id then + logger.log("no contest set. run :CP first", vim.log.levels.ERROR) return end - config = config_module.setup(user_config) + cache.load() + local contest_data = cache.get_contest_data(vim.g.cp.platform, vim.g.cp.contest_id) + if not contest_data or not contest_data.problems then + logger.log("no contest metadata found. set up a problem first to cache contest data", vim.log.levels.ERROR) + return + end + + local problems = contest_data.problems + local current_problem_id + + if vim.g.cp.platform == "cses" then + current_problem_id = vim.g.cp.contest_id + else + current_problem_id = vim.g.cp.problem_id + end + + if not current_problem_id then + logger.log("no current problem set", vim.log.levels.ERROR) + return + end + + local current_index = nil + 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 + logger.log("current problem not found in contest", vim.log.levels.ERROR) + return + end + + local new_index = current_index + delta + + if new_index < 1 or new_index > #problems then + local direction = delta > 0 and "next" or "previous" + logger.log(("no %s problem available"):format(direction), vim.log.levels.INFO) + return + end + + local new_problem = problems[new_index] + + if vim.g.cp.platform == "cses" then + setup_problem(new_problem.id) + else + setup_problem(vim.g.cp.contest_id, new_problem.id) + end +end + +local function ensure_initialized() + if config then + return + end + + if not lazy_require() then + return + end + + vim.g.cp = vim.g.cp or {} + config = config_module.setup(vim.g.cp.config) logger.set_config(config) snippets.setup(config) - initialized = true +end + +local function parse_command(args) + if #args == 0 then + return { type = "error", message = "Usage: :CP [problem] | :CP | :CP " } + end + + local first = args[1] + + if vim.tbl_contains(actions, first) then + return { type = "action", action = first } + end + + if vim.tbl_contains(platforms, first) then + if #args == 1 then + return { type = "platform_only", platform = first } + elseif #args == 2 then + return { type = "contest_setup", platform = first, contest = args[2] } + elseif #args == 3 then + return { type = "full_setup", platform = first, contest = args[2], problem = args[3] } + else + return { type = "error", message = "Too many arguments" } + end + end + + if vim.g.cp and vim.g.cp.platform and vim.g.cp.contest_id then + return { type = "problem_switch", problem = first } + end + + return { type = "error", message = "Unknown command or no contest context" } end function M.handle_command(opts) - local args = opts.fargs - if #args == 0 then - logger.log("Usage: :CP ", vim.log.levels.ERROR) + ensure_initialized() + local cmd = parse_command(opts.fargs) + + if cmd.type == "error" then + logger.log(cmd.message, vim.log.levels.ERROR) return end - local cmd = args[1] + if cmd.type == "action" then + if cmd.action == "run" then + run_problem() + elseif cmd.action == "debug" then + debug_problem() + elseif cmd.action == "diff" then + diff_problem() + elseif cmd.action == "next" then + navigate_problem(1) + elseif cmd.action == "prev" then + navigate_problem(-1) + end + return + end - if vim.tbl_contains(competition_types, cmd) then - if args[2] then - setup_contest(cmd) - if (cmd == "atcoder" or cmd == "codeforces") and args[3] then - setup_problem(args[2], args[3]) + if cmd.type == "platform_only" then + set_platform(cmd.platform) + return + end + + if cmd.type == "contest_setup" then + if set_platform(cmd.platform) then + vim.g.cp.contest_id = cmd.contest + local metadata_result = scrape.scrape_contest_metadata(cmd.platform, cmd.contest) + if not metadata_result.success then + logger.log( + "failed to load contest metadata: " .. (metadata_result.error or "unknown error"), + vim.log.levels.WARN + ) else - setup_problem(args[2]) + logger.log( + ("loaded %d problems for %s %s"):format(#metadata_result.problems, cmd.platform, cmd.contest) + ) end - else - setup_contest(cmd) end - elseif cmd == "run" then - run_problem() - elseif cmd == "debug" then - debug_problem() - elseif cmd == "diff" then - diff_problem() - elseif vim.g.cp_contest and not vim.tbl_contains(competition_types, cmd) then - if (vim.g.cp_contest == "atcoder" or vim.g.cp_contest == "codeforces") and args[2] then - setup_problem(cmd, args[2]) - else - setup_problem(cmd) + return + end + + if cmd.type == "full_setup" then + if set_platform(cmd.platform) then + vim.g.cp.contest_id = cmd.contest + setup_problem(cmd.contest, cmd.problem) end - else - logger.log( - ("unknown contest type '%s'. Available: [%s]"):format(cmd, table.concat(competition_types, ", ")), - vim.log.levels.ERROR - ) + return + end + + if cmd.type == "problem_switch" then + if vim.g.cp.platform == "cses" then + setup_problem(cmd.problem) + else + setup_problem(vim.g.cp.contest_id, cmd.problem) + end + return end end +function M.is_initialized() + return config ~= nil +end + return M diff --git a/lua/cp/problem.lua b/lua/cp/problem.lua index 2ae83a7..4aa73d2 100644 --- a/lua/cp/problem.lua +++ b/lua/cp/problem.lua @@ -3,7 +3,7 @@ ---@field contest_id string Contest ID (e.g. "abc123", "1933") ---@field problem_id? string Problem ID for AtCoder/Codeforces (e.g. "a", "b") ---@field source_file string Source filename (e.g. "abc123a.cpp") ----@field binary_file string Binary output path (e.g. "build/abc123a") +---@field binary_file string Binary output path (e.g. "build/abc123a.run") ---@field input_file string Input test file path (e.g. "io/abc123a.in") ---@field output_file string Output file path (e.g. "io/abc123a.out") ---@field expected_file string Expected output path (e.g. "io/abc123a.expected") @@ -26,7 +26,7 @@ function M.create_context(contest, contest_id, problem_id, config) contest_id = contest_id, problem_id = problem_id, source_file = source_file, - binary_file = ("build/%s"):format(base_name), + binary_file = ("build/%s.run"):format(base_name), input_file = ("io/%s.in"):format(base_name), output_file = ("io/%s.out"):format(base_name), expected_file = ("io/%s.expected"):format(base_name), diff --git a/lua/cp/scrape.lua b/lua/cp/scrape.lua index b03a5c1..c2436e4 100644 --- a/lua/cp/scrape.lua +++ b/lua/cp/scrape.lua @@ -1,5 +1,6 @@ local M = {} local logger = require("cp.log") +local cache = require("cp.cache") local function get_plugin_path() local plugin_path = debug.getinfo(1, "S").source:sub(2) @@ -35,7 +36,72 @@ local function setup_python_env() return true end +---@param platform string +---@param contest_id string +---@return {success: boolean, problems?: table[], error?: string} +function M.scrape_contest_metadata(platform, contest_id) + cache.load() + + local cached_data = cache.get_contest_data(platform, contest_id) + if cached_data then + return { + success = true, + problems = cached_data.problems, + } + end + + if not setup_python_env() then + return { + success = false, + error = "Python environment setup failed", + } + end + + local plugin_path = get_plugin_path() + local scraper_path = plugin_path .. "/scrapers/" .. platform .. ".py" + local args = { "uv", "run", scraper_path, "metadata", contest_id } + + local result = vim.system(args, { + cwd = plugin_path, + text = true, + timeout = 30000, + }):wait() + + if result.code ~= 0 then + return { + success = false, + error = "Failed to run metadata scraper: " .. (result.stderr or "Unknown error"), + } + end + + local ok, data = pcall(vim.json.decode, result.stdout) + if not ok then + return { + success = false, + error = "Failed to parse metadata scraper output: " .. tostring(data), + } + end + + if not data.success then + return data + end + + local problems_list + if platform == "cses" then + problems_list = data.categories and data.categories["CSES Problem Set"] or {} + else + problems_list = data.problems or {} + end + + cache.set_contest_data(platform, contest_id, problems_list) + return { + success = true, + problems = problems_list, + } +end + ---@param ctx ProblemContext +---@return {success: boolean, problem_id: string, test_count?: number, url?: string, error?: string} function M.scrape_problem(ctx) ensure_io_directory() @@ -50,6 +116,7 @@ function M.scrape_problem(ctx) if not setup_python_env() then return { success = false, + problem_id = ctx.problem_name, error = "Python environment setup failed", } end @@ -59,9 +126,9 @@ function M.scrape_problem(ctx) local args if ctx.contest == "cses" then - args = { "uv", "run", scraper_path, ctx.contest_id } + args = { "uv", "run", scraper_path, "tests", ctx.contest_id } else - args = { "uv", "run", scraper_path, ctx.contest_id, ctx.problem_id } + args = { "uv", "run", scraper_path, "tests", ctx.contest_id, ctx.problem_id } end local result = vim.system(args, { @@ -73,7 +140,8 @@ function M.scrape_problem(ctx) if result.code ~= 0 then return { success = false, - error = "Failed to run scraper: " .. (result.stderr or "Unknown error"), + problem_id = ctx.problem_name, + error = "Failed to run tests scraper: " .. (result.stderr or "Unknown error"), } end @@ -81,7 +149,8 @@ function M.scrape_problem(ctx) if not ok then return { success = false, - error = "Failed to parse scraper output: " .. tostring(data), + problem_id = ctx.problem_name, + error = "Failed to parse tests scraper output: " .. tostring(data), } end @@ -89,7 +158,7 @@ function M.scrape_problem(ctx) return data end - if #data.test_cases > 0 then + if data.test_cases and #data.test_cases > 0 then local all_inputs = {} local all_outputs = {} @@ -113,7 +182,7 @@ function M.scrape_problem(ctx) return { success = true, problem_id = ctx.problem_name, - test_count = #data.test_cases, + test_count = data.test_cases and #data.test_cases or 0, url = data.url, } end diff --git a/plugin/cp.lua b/plugin/cp.lua index 9fc7e28..690823a 100644 --- a/plugin/cp.lua +++ b/plugin/cp.lua @@ -3,20 +3,54 @@ if vim.g.loaded_cp then end vim.g.loaded_cp = 1 -local competition_types = { "atcoder", "codeforces", "cses" } +local platforms = { "atcoder", "codeforces", "cses" } +local actions = { "run", "debug", "diff", "next", "prev" } vim.api.nvim_create_user_command("CP", function(opts) local cp = require("cp") - if not cp.is_initialized() then - cp.setup() - end cp.handle_command(opts) end, { nargs = "*", - complete = function(ArgLead, _, _) - local commands = vim.list_extend(vim.deepcopy(competition_types), { "run", "debug", "diff" }) - return vim.tbl_filter(function(cmd) - return cmd:find(ArgLead, 1, true) == 1 - end, commands) + complete = function(ArgLead, CmdLine, _) + local args = vim.split(vim.trim(CmdLine), "%s+") + local num_args = #args + if CmdLine:sub(-1) == " " then + num_args = num_args + 1 + end + + if num_args == 2 then + local candidates = {} + vim.list_extend(candidates, platforms) + vim.list_extend(candidates, actions) + if vim.g.cp and vim.g.cp.platform and vim.g.cp.contest_id then + local cache = require("cp.cache") + cache.load() + local contest_data = cache.get_contest_data(vim.g.cp.platform, vim.g.cp.contest_id) + if contest_data and contest_data.problems then + for _, problem in ipairs(contest_data.problems) do + table.insert(candidates, problem.id) + end + end + end + return vim.tbl_filter(function(cmd) + return cmd:find(ArgLead, 1, true) == 1 + end, candidates) + elseif num_args == 4 then + if vim.tbl_contains(platforms, args[2]) then + local cache = require("cp.cache") + cache.load() + local contest_data = cache.get_contest_data(args[2], args[3]) + if contest_data and contest_data.problems then + local candidates = {} + for _, problem in ipairs(contest_data.problems) do + table.insert(candidates, problem.id) + end + return vim.tbl_filter(function(cmd) + return cmd:find(ArgLead, 1, true) == 1 + end, candidates) + end + end + end + return {} end, }) diff --git a/scrapers/atcoder.py b/scrapers/atcoder.py index 1f84f2e..63bd3db 100644 --- a/scrapers/atcoder.py +++ b/scrapers/atcoder.py @@ -8,13 +8,58 @@ from bs4 import BeautifulSoup def parse_problem_url(contest_id: str, problem_letter: str) -> str: - task_id = f"{contest_id}_{problem_letter}" + task_id: str = f"{contest_id}_{problem_letter}" return f"https://atcoder.jp/contests/{contest_id}/tasks/{task_id}" +def scrape_contest_problems(contest_id: str) -> list[dict[str, str]]: + try: + contest_url: str = f"https://atcoder.jp/contests/{contest_id}/tasks" + headers: dict[str, str] = { + "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" + } + + response = requests.get(contest_url, headers=headers, timeout=10) + response.raise_for_status() + + soup = BeautifulSoup(response.text, "html.parser") + problems: list[dict[str, str]] = [] + + task_table = soup.find("table", class_="table") + if not task_table: + return [] + + rows = task_table.find_all("tr")[1:] # Skip header row + + for row in rows: + cells = row.find_all("td") + if len(cells) >= 2: + task_link = cells[1].find("a") + if task_link: + task_name: str = task_link.get_text(strip=True) + task_href: str = task_link.get("href", "") + + # Extract problem letter from task name or URL + task_id: str = task_href.split("/")[-1] if task_href else "" + if task_id.startswith(contest_id + "_"): + problem_letter: str = task_id[len(contest_id) + 1 :] + + if problem_letter and task_name: + problems.append( + {"id": problem_letter.lower(), "name": task_name} + ) + + problems.sort(key=lambda x: x["id"]) + return problems + + except Exception as e: + print(f"Failed to scrape AtCoder contest problems: {e}", file=sys.stderr) + return [] + + def scrape(url: str) -> list[tuple[str, str]]: try: - headers = { + headers: dict[str, str] = { "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" } @@ -23,7 +68,7 @@ def scrape(url: str) -> list[tuple[str, str]]: soup = BeautifulSoup(response.text, "html.parser") - tests = [] + tests: list[tuple[str, str]] = [] sample_headers = soup.find_all( "h3", string=lambda x: x and "sample" in x.lower() if x else False @@ -39,8 +84,10 @@ def scrape(url: str) -> list[tuple[str, str]]: if "output" in next_header.get_text().lower(): output_pre = next_header.find_next("pre") if output_pre: - input_text = input_pre.get_text().strip().replace("\r", "") - output_text = ( + input_text: str = ( + input_pre.get_text().strip().replace("\r", "") + ) + output_text: str = ( output_pre.get_text().strip().replace("\r", "") ) if input_text and output_text: @@ -56,55 +103,101 @@ def scrape(url: str) -> list[tuple[str, str]]: return [] -def main(): - if len(sys.argv) != 3: - result = { +def main() -> None: + if len(sys.argv) < 2: + result: dict[str, str | bool] = { "success": False, - "error": "Usage: atcoder.py ", - "problem_id": None, + "error": "Usage: atcoder.py metadata OR atcoder.py tests ", } print(json.dumps(result)) sys.exit(1) - contest_id = sys.argv[1] - problem_letter = sys.argv[2] - problem_id = contest_id + problem_letter.lower() + mode: str = sys.argv[1] - url = parse_problem_url(contest_id, problem_letter) - print(f"Scraping: {url}", file=sys.stderr) + if mode == "metadata": + if len(sys.argv) != 3: + result = { + "success": False, + "error": "Usage: atcoder.py metadata ", + } + print(json.dumps(result)) + sys.exit(1) - tests = scrape(url) + contest_id: str = sys.argv[2] + problems: list[dict[str, str]] = scrape_contest_problems(contest_id) + + if not problems: + result = { + "success": False, + "error": f"No problems found for contest {contest_id}", + } + print(json.dumps(result)) + sys.exit(1) - if not tests: result = { - "success": False, - "error": f"No tests found for {contest_id} {problem_letter}", + "success": True, + "contest_id": contest_id, + "problems": problems, + } + print(json.dumps(result)) + + elif mode == "tests": + if len(sys.argv) != 4: + result = { + "success": False, + "error": "Usage: atcoder.py tests ", + } + print(json.dumps(result)) + sys.exit(1) + + contest_id: str = sys.argv[2] + problem_letter: str = sys.argv[3] + problem_id: str = contest_id + problem_letter.lower() + + url: str = parse_problem_url(contest_id, problem_letter) + print(f"Scraping: {url}", file=sys.stderr) + + tests: list[tuple[str, str]] = scrape(url) + + if not tests: + result = { + "success": False, + "error": f"No tests found for {contest_id} {problem_letter}", + "problem_id": problem_id, + "url": url, + } + print(json.dumps(result)) + sys.exit(1) + + test_cases: list[dict[str, str]] = [] + for input_data, output_data in tests: + test_cases.append({"input": input_data, "output": output_data}) + + if test_cases: + combined_input: str = ( + str(len(test_cases)) + + "\n" + + "\n".join(tc["input"] for tc in test_cases) + ) + combined_output: str = "\n".join(tc["output"] for tc in test_cases) + test_cases = [{"input": combined_input, "output": combined_output}] + + result = { + "success": True, "problem_id": problem_id, "url": url, + "test_cases": test_cases, + } + print(json.dumps(result)) + + else: + result = { + "success": False, + "error": f"Unknown mode: {mode}. Use 'metadata' or 'tests'", } print(json.dumps(result)) sys.exit(1) - test_cases = [] - for input_data, output_data in tests: - test_cases.append({"input": input_data, "output": output_data}) - - if test_cases: - combined_input = ( - str(len(test_cases)) + "\n" + "\n".join(tc["input"] for tc in test_cases) - ) - combined_output = "\n".join(tc["output"] for tc in test_cases) - test_cases = [{"input": combined_input, "output": combined_output}] - - result = { - "success": True, - "problem_id": problem_id, - "url": url, - "test_cases": test_cases, - } - - print(json.dumps(result)) - if __name__ == "__main__": main() diff --git a/scrapers/codeforces.py b/scrapers/codeforces.py index 35589a0..73be5a4 100644 --- a/scrapers/codeforces.py +++ b/scrapers/codeforces.py @@ -7,14 +7,14 @@ import cloudscraper from bs4 import BeautifulSoup -def scrape(url: str): +def scrape(url: str) -> list[tuple[str, str]]: try: scraper = cloudscraper.create_scraper() response = scraper.get(url, timeout=10) response.raise_for_status() soup = BeautifulSoup(response.text, "html.parser") - tests = [] + tests: list[tuple[str, str]] = [] input_sections = soup.find_all("div", class_="input") output_sections = soup.find_all("div", class_="output") @@ -24,8 +24,8 @@ def scrape(url: str): out_pre = out_section.find("pre") if inp_pre and out_pre: - input_lines = [] - output_lines = [] + input_lines: list[str] = [] + output_lines: list[str] = [] for line_div in inp_pre.find_all("div", class_="test-example-line"): input_lines.append(line_div.get_text().strip()) @@ -60,51 +60,134 @@ def parse_problem_url(contest_id: str, problem_letter: str) -> str: ) -def scrape_sample_tests(url: str): +def scrape_contest_problems(contest_id: str) -> list[dict[str, str]]: + try: + contest_url: str = f"https://codeforces.com/contest/{contest_id}" + scraper = cloudscraper.create_scraper() + response = scraper.get(contest_url, timeout=10) + response.raise_for_status() + + soup = BeautifulSoup(response.text, "html.parser") + problems: list[dict[str, str]] = [] + + problem_links = soup.find_all( + "a", href=lambda x: x and f"/contest/{contest_id}/problem/" in x + ) + + for link in problem_links: + href: str = link.get("href", "") + if f"/contest/{contest_id}/problem/" in href: + problem_letter: str = href.split("/")[-1].lower() + problem_name: str = link.get_text(strip=True) + + if problem_letter and problem_name and len(problem_letter) == 1: + problems.append({"id": problem_letter, "name": problem_name}) + + problems.sort(key=lambda x: x["id"]) + + seen: set[str] = set() + unique_problems: list[dict[str, str]] = [] + for p in problems: + if p["id"] not in seen: + seen.add(p["id"]) + unique_problems.append(p) + + return unique_problems + + except Exception as e: + print(f"Failed to scrape contest problems: {e}", file=sys.stderr) + return [] + + +def scrape_sample_tests(url: str) -> list[tuple[str, str]]: print(f"Scraping: {url}", file=sys.stderr) return scrape(url) -def main(): - if len(sys.argv) != 3: - result = { +def main() -> None: + if len(sys.argv) < 2: + result: dict[str, str | bool] = { "success": False, - "error": "Usage: codeforces.py ", - "problem_id": None, + "error": "Usage: codeforces.py metadata OR codeforces.py tests ", } print(json.dumps(result)) sys.exit(1) - contest_id = sys.argv[1] - problem_letter = sys.argv[2] - problem_id = contest_id + problem_letter.lower() + mode: str = sys.argv[1] - url = parse_problem_url(contest_id, problem_letter) - tests = scrape_sample_tests(url) + if mode == "metadata": + if len(sys.argv) != 3: + result = { + "success": False, + "error": "Usage: codeforces.py metadata ", + } + print(json.dumps(result)) + sys.exit(1) + + contest_id: str = sys.argv[2] + problems: list[dict[str, str]] = scrape_contest_problems(contest_id) + + if not problems: + result = { + "success": False, + "error": f"No problems found for contest {contest_id}", + } + print(json.dumps(result)) + sys.exit(1) - if not tests: result = { - "success": False, - "error": f"No tests found for {contest_id} {problem_letter}", + "success": True, + "contest_id": contest_id, + "problems": problems, + } + print(json.dumps(result)) + + elif mode == "tests": + if len(sys.argv) != 4: + result = { + "success": False, + "error": "Usage: codeforces.py tests ", + } + print(json.dumps(result)) + sys.exit(1) + + contest_id: str = sys.argv[2] + problem_letter: str = sys.argv[3] + problem_id: str = contest_id + problem_letter.lower() + + url: str = parse_problem_url(contest_id, problem_letter) + tests: list[tuple[str, str]] = scrape_sample_tests(url) + + if not tests: + result = { + "success": False, + "error": f"No tests found for {contest_id} {problem_letter}", + "problem_id": problem_id, + "url": url, + } + print(json.dumps(result)) + sys.exit(1) + + test_cases: list[dict[str, str]] = [] + for input_data, output_data in tests: + test_cases.append({"input": input_data, "output": output_data}) + + result = { + "success": True, "problem_id": problem_id, "url": url, + "test_cases": test_cases, + } + print(json.dumps(result)) + + else: + result = { + "success": False, + "error": f"Unknown mode: {mode}. Use 'metadata' or 'tests'", } print(json.dumps(result)) sys.exit(1) - test_cases = [] - for input_data, output_data in tests: - test_cases.append({"input": input_data, "output": output_data}) - - result = { - "success": True, - "problem_id": problem_id, - "url": url, - "test_cases": test_cases, - } - - print(json.dumps(result)) - if __name__ == "__main__": main() diff --git a/scrapers/cses.py b/scrapers/cses.py index 8da2ba6..8cd6020 100755 --- a/scrapers/cses.py +++ b/scrapers/cses.py @@ -15,9 +15,54 @@ def parse_problem_url(problem_input: str) -> str | None: return None +def scrape_all_problems() -> dict[str, list[dict[str, str]]]: + try: + problemset_url: str = "https://cses.fi/problemset/" + headers: dict[str, str] = { + "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" + } + + response = requests.get(problemset_url, headers=headers, timeout=10) + response.raise_for_status() + + soup = BeautifulSoup(response.text, "html.parser") + all_categories: dict[str, list[dict[str, str]]] = {} + + problem_links = soup.find_all( + "a", href=lambda x: x and "/problemset/task/" in x + ) + print(f"Found {len(problem_links)} problem links", file=sys.stderr) + + current_category: str | None = None + for element in soup.find_all(["h1", "a"]): + if element.name == "h1": + current_category = element.get_text().strip() + if current_category not in all_categories: + all_categories[current_category] = [] + elif element.name == "a" and "/problemset/task/" in element.get("href", ""): + href: str = element.get("href", "") + problem_id: str = href.split("/")[-1] + problem_name: str = element.get_text(strip=True) + + if problem_id.isdigit() and problem_name and current_category: + all_categories[current_category].append( + {"id": problem_id, "name": problem_name} + ) + + for category in all_categories: + all_categories[category].sort(key=lambda x: int(x["id"])) + + print(f"Found {len(all_categories)} categories", file=sys.stderr) + return all_categories + + except Exception as e: + print(f"Failed to scrape CSES problems: {e}", file=sys.stderr) + return {} + + def scrape(url: str) -> list[tuple[str, str]]: try: - headers = { + headers: dict[str, str] = { "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" } @@ -26,13 +71,13 @@ def scrape(url: str) -> list[tuple[str, str]]: soup = BeautifulSoup(response.text, "html.parser") - tests = [] + tests: list[tuple[str, str]] = [] example_header = soup.find("h1", string="Example") if example_header: current = example_header.find_next_sibling() - input_text = None - output_text = None + input_text: str | None = None + output_text: str | None = None while current: if current.name == "p" and "Input:" in current.get_text(): @@ -56,57 +101,99 @@ def scrape(url: str) -> list[tuple[str, str]]: return [] -def main(): - if len(sys.argv) != 2: - result = { +def main() -> None: + if len(sys.argv) < 2: + result: dict[str, str | bool] = { "success": False, - "error": "Usage: cses.py ", - "problem_id": None, + "error": "Usage: cses.py metadata OR cses.py tests ", } print(json.dumps(result)) sys.exit(1) - problem_input = sys.argv[1] - url = parse_problem_url(problem_input) + mode: str = sys.argv[1] + + if mode == "metadata": + if len(sys.argv) != 2: + result = { + "success": False, + "error": "Usage: cses.py metadata", + } + print(json.dumps(result)) + sys.exit(1) + + all_categories: dict[str, list[dict[str, str]]] = scrape_all_problems() + + if not all_categories: + result = { + "success": False, + "error": "Failed to scrape CSES problem categories", + } + print(json.dumps(result)) + sys.exit(1) - if not url: result = { - "success": False, - "error": f"Invalid problem input: {problem_input}. Use either problem ID (e.g., 1068) or full URL", - "problem_id": problem_input if problem_input.isdigit() else None, + "success": True, + "categories": all_categories, } print(json.dumps(result)) - sys.exit(1) - tests = scrape(url) + elif mode == "tests": + if len(sys.argv) != 3: + result = { + "success": False, + "error": "Usage: cses.py tests ", + } + print(json.dumps(result)) + sys.exit(1) - problem_id = ( - problem_input if problem_input.isdigit() else problem_input.split("/")[-1] - ) + problem_input: str = sys.argv[2] + url: str | None = parse_problem_url(problem_input) + + if not url: + result = { + "success": False, + "error": f"Invalid problem input: {problem_input}. Use either problem ID (e.g., 1068) or full URL", + "problem_id": problem_input if problem_input.isdigit() else None, + } + print(json.dumps(result)) + sys.exit(1) + + tests: list[tuple[str, str]] = scrape(url) + + problem_id: str = ( + problem_input if problem_input.isdigit() else problem_input.split("/")[-1] + ) + + if not tests: + result = { + "success": False, + "error": f"No tests found for {problem_input}", + "problem_id": problem_id, + "url": url, + } + print(json.dumps(result)) + sys.exit(1) + + test_cases: list[dict[str, str]] = [] + for input_data, output_data in tests: + test_cases.append({"input": input_data, "output": output_data}) - if not tests: result = { - "success": False, - "error": f"No tests found for {problem_input}", + "success": True, "problem_id": problem_id, "url": url, + "test_cases": test_cases, + } + print(json.dumps(result)) + + else: + result = { + "success": False, + "error": f"Unknown mode: {mode}. Use 'metadata' or 'tests'", } print(json.dumps(result)) sys.exit(1) - test_cases = [] - for input_data, output_data in tests: - test_cases.append({"input": input_data, "output": output_data}) - - result = { - "success": True, - "problem_id": problem_id, - "url": url, - "test_cases": test_cases, - } - - print(json.dumps(result)) - if __name__ == "__main__": main()