feat: flesh out language support

after/ftplugin/cpout.lua

@@ -4,4 +4,4 @@ vim.opt_local.statuscolumn = ""
 vim.opt_local.signcolumn = "no"
 vim.opt_local.wrap = true
 vim.opt_local.linebreak = true
-vim.opt_local.modifiable = false
+vim.opt_local.modifiable = true

doc/cp.txt

@@ -27,11 +27,12 @@ cp.nvim uses a single :CP command with intelligent argument parsing:
 
 Setup Commands ~
 
-:CP {platform} {contest_id} {problem_id}
+:CP {platform} {contest_id} {problem_id} [--lang={language}]
                                 Full setup: set platform, load contest metadata,
                                 and set up specific problem. Scrapes test cases
                                 and creates source file.
                                 Example: :CP codeforces 1933 a
+                                Example: :CP codeforces 1933 a --lang=python
 
 :CP {platform} {contest_id}     Contest setup: set platform and load contest
                                 metadata for navigation. Caches problem list.
@@ -40,9 +41,11 @@ Setup Commands ~
 :CP {platform}                  Platform setup: set platform only.
                                 Example: :CP cses
 
-:CP {problem_id}                Problem switch: switch to different problem
+:CP {problem_id} [--lang={language}]
+                                Problem switch: switch to different problem
                                 within current contest context.
                                 Example: :CP b (switch to problem b)
+                                Example: :CP b --lang=python
 
 Action Commands ~
 
@@ -75,21 +78,36 @@ Optional configuration with lazy.nvim: >
           debug = false,
           contests = {
             default = {
-              cpp_version = 20,
+              cpp = {
-              compile_flags = { "-O2", "-DLOCAL", "-Wall", "-Wextra" },
+                compile = {
-              debug_flags = { "-g3", "-fsanitize=address,undefined", "-DLOCAL" },
+                  'g++', '-std=c++{version}', '-O2', '-Wall', '-Wextra',
+                  '-DLOCAL', '{source}', '-o', '{binary}',
+                },
+                run = { '{binary}' },
+                debug = {
+                  'g++', '-std=c++{version}', '-g3',
+                  '-fsanitize=address,undefined', '-DLOCAL',
+                  '{source}', '-o', '{binary}',
+                },
+                version = 20,
+                extension = "cc",
+              },
+              python = {
+                run = { 'python3', '{source}' },
+                debug = { 'python3', '{source}' },
+                extension = "py",
+              },
               timeout_ms = 2000,
             },
-            atcoder = { cpp_version = 23 },
+            codeforces = { cpp = { version = 23 } },
           },
           hooks = {
             before_run = function(problem_id) vim.cmd.w() end,
             before_debug = function(problem_id) ... end,
           },
-          tile = function(source_buf, input_buf, output_buf)
+          snippets = { ... },  -- LuaSnip snippets
-          end,
+          tile = function(source_buf, input_buf, output_buf) ... end,
-          filename = function(contest, problem_id, problem_letter)
+          filename = function(contest, problem_id, problem_letter) ... end,
-          end,
         }
     }
 <
@@ -98,10 +116,21 @@ Configuration options:
 
 contests                        Dictionary of contest configurations - each contest inherits from 'default'.
 
-  cpp_version                   c++ standard version (e.g. 20, 23)
+  cpp                           C++ language configuration
-  compile_flags                 compiler flags for run builds
+    compile                     Compile command template with {version}, {source}, {binary} placeholders
-  debug_flags                   compiler flags for debug builds
+    run                         Run command template with {binary} placeholder
-  timeout_ms                    duration (ms) to run/debug before timeout
+    debug                       Debug compile command template
+    version                     C++ standard version (e.g. 20, 23)
+    extension                   File extension for C++ files (default: "cc")
+
+  python                        Python language configuration
+    run                         Run command template with {source} placeholder
+    debug                       Debug run command template
+    extension                   File extension for Python files (default: "py")
+
+  default_language              Default language when --lang not specified (default: "cpp")
+
+  timeout_ms                    Duration (ms) to run/debug before timeout
 
 snippets                        LuaSnip snippets by contest type
 
@@ -221,8 +250,8 @@ cp.nvim creates the following file structure upon problem setup:
     build/
       {contest_id}{problem_id}.run  " Compiled binary
     io/
-      {contest_id}{problem_id}.in       " Test input
+      {contest_id}{problem_id}.cpin     " Test input
-      {contest_id}{problem_id}.out      " Program output
+      {contest_id}{problem_id}.cpout    " Program output
       {contest_id}{problem_id}.expected " Expected output
 
 The plugin automatically manages this structure and navigation between problems
@@ -233,9 +262,17 @@ SNIPPETS						*cp-snippets*
 cp.nvim integrates with LuaSnip for automatic template expansion. When you
 open a new problem file, type the contest name and press <Tab> to expand.
 
-Built-in snippets include basic C++ templates for each contest type.
+Built-in snippets include basic C++ and Python templates for each contest type.
 Custom snippets can be added via configuration.
 
+IMPORTANT: Snippet trigger names must exactly match the contest/platform names:
+- "codeforces" for Codeforces problems
+- "atcoder" for AtCoder problems
+- "cses" for CSES problems
+
+The plugin automatically selects the appropriate template based on the file
+extension (e.g., .cc files get C++ templates, .py files get Python templates).
+
 HEALTH CHECK						*cp-health*
 
 Run |:checkhealth| cp to verify your setup.

lua/cp/config.lua

@@ -1,13 +1,34 @@
+---@class LanguageConfig
+---@field compile? string[] Compile command template
+---@field run string[] Run command template
+---@field debug? string[] Debug command template
+---@field executable? string Executable name
+---@field version? number Language version
+---@field extension string File extension
+
+---@class ContestConfig
+---@field cpp LanguageConfig
+---@field python LanguageConfig
+---@field default_language string
+---@field timeout_ms number
+
 ---@class cp.Config
----@field contests table
+---@field contests table<string, ContestConfig>
----@field snippets table
+---@field snippets table[]
 ---@field hooks table
 ---@field debug boolean
 ---@field tile? fun(source_buf: number, input_buf: number, output_buf: number)
----@field filename? fun(contest: string, problem_id: string, problem_letter?: string): string
+---@field filename? fun(contest: string, contest_id: string, problem_id?: string, config: cp.Config, language?: string): string
 
 local M = {}
 
+local filetype_to_language = {
+	cc = "cpp",
+	c = "cpp",
+	py = "python",
+	py3 = "python",
+}
+
 ---@type cp.Config
 M.defaults = {
 	contests = {
@@ -37,13 +58,16 @@ M.defaults = {
 				},
 				executable = nil,
 				version = 20,
+				extension = "cc",
 			},
 			python = {
 				compile = nil,
 				run = { "{source}" },
 				debug = { "{source}" },
 				executable = "python3",
+				extension = "py",
 			},
+			default_language = "cpp",
 			timeout_ms = 2000,
 		},
 		atcoder = {
@@ -74,8 +98,8 @@ local function extend_contest_config(base_config, contest_config)
 	return result
 end
 
----@param user_config table|nil
+---@param user_config cp.Config|nil
----@return table
+---@return cp.Config
 function M.setup(user_config)
 	vim.validate({
 		user_config = { user_config, { "table", "nil" }, true },
@@ -97,6 +121,20 @@ function M.setup(user_config)
 				before_debug = { user_config.hooks.before_debug, { "function", "nil" }, true },
 			})
 		end
+
+		if user_config.contests then
+			for contest_name, contest_config in pairs(user_config.contests) do
+				for lang_name, lang_config in pairs(contest_config) do
+					if type(lang_config) == "table" and lang_config.extension then
+						if not vim.tbl_contains(vim.tbl_keys(filetype_to_language), lang_config.extension) then
+							error(("Invalid extension '%s' for language '%s' in contest '%s'. Valid extensions: %s"):format(
+								lang_config.extension, lang_name, contest_name, table.concat(vim.tbl_keys(filetype_to_language), ", ")
+							))
+						end
+					end
+				end
+			end
+		end
 	end
 
 	local config = vim.tbl_deep_extend("force", M.defaults, user_config or {})
@@ -111,14 +149,32 @@ function M.setup(user_config)
 	return config
 end
 
-local function default_filename(contest, contest_id, problem_id)
+---@param contest string
+---@param contest_id string
+---@param problem_id? string
+---@param config cp.Config
+---@param language? string
+---@return string
+local function default_filename(contest, contest_id, problem_id, config, language)
+	vim.validate({
+		contest = { contest, "string" },
+		contest_id = { contest_id, "string" },
+		problem_id = { problem_id, { "string", "nil" }, true },
+		config = { config, "table" },
+		language = { language, { "string", "nil" }, true },
+	})
+
 	local full_problem_id = contest_id:lower()
 	if contest == "atcoder" or contest == "codeforces" then
 		if problem_id then
 			full_problem_id = full_problem_id .. problem_id:lower()
 		end
 	end
-	return full_problem_id .. ".cc"
+
+	local contest_config = config.contests[contest] or config.contests.default
+	local target_language = language or contest_config.default_language
+	local language_config = contest_config[target_language]
+	return full_problem_id .. "." .. language_config.extension
 end
 
 M.default_filename = default_filename

lua/cp/init.lua

@@ -48,7 +48,8 @@ end
 
 ---@param contest_id string
 ---@param problem_id? string
-local function setup_problem(contest_id, problem_id)
+---@param language? string
+local function setup_problem(contest_id, problem_id, language)
 	if not state.platform then
 		logger.log("no platform set. run :CP <platform> <contest> first", vim.log.levels.ERROR)
 		return
@@ -88,7 +89,7 @@ local function setup_problem(contest_id, problem_id)
 		state.test_cases = cached_test_cases
 	end
 
-	local ctx = problem.create_context(state.platform, contest_id, problem_id, config)
+	local ctx = problem.create_context(state.platform, contest_id, problem_id, config, language)
 
 	local scrape_result = scrape.scrape_problem(ctx)
 
@@ -305,33 +306,52 @@ end
 
 local function parse_command(args)
 	if #args == 0 then
-		return { type = "error", message = "Usage: :CP <platform> <contest> [problem] | :CP <action> | :CP <problem>" }
+		return { type = "error", message = "Usage: :CP <platform> <contest> [problem] [--lang=<language>] | :CP <action> | :CP <problem>" }
 	end
 
-	local first = args[1]
+	local language = nil
+
+	for i, arg in ipairs(args) do
+		local lang_match = arg:match("^--lang=(.+)$")
+		if lang_match then
+			language = lang_match
+		elseif arg == "--lang" then
+			if i + 1 <= #args then
+				language = args[i + 1]
+			else
+				return { type = "error", message = "--lang requires a value" }
+			end
+		end
+	end
+
+	local filtered_args = vim.tbl_filter(function(arg)
+		return not (arg:match("^--lang") or arg == language)
+	end, args)
+
+	local first = filtered_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
+		if #filtered_args == 1 then
-			return { type = "platform_only", platform = first }
+			return { type = "platform_only", platform = first, language = language }
-		elseif #args == 2 then
+		elseif #filtered_args == 2 then
 			if first == "cses" then
-				return { type = "cses_problem", platform = first, problem = args[2] }
+				return { type = "cses_problem", platform = first, problem = filtered_args[2], language = language }
 			else
-				return { type = "contest_setup", platform = first, contest = args[2] }
+				return { type = "contest_setup", platform = first, contest = filtered_args[2], language = language }
 			end
-		elseif #args == 3 then
+		elseif #filtered_args == 3 then
-			return { type = "full_setup", platform = first, contest = args[2], problem = args[3] }
+			return { type = "full_setup", platform = first, contest = filtered_args[2], problem = filtered_args[3], language = language }
 		else
 			return { type = "error", message = "Too many arguments" }
 		end
 	end
 
 	if state.platform and state.contest_id then
-		return { type = "problem_switch", problem = first }
+		return { type = "problem_switch", problem = first, language = language }
 	end
 
 	return { type = "error", message = "Unknown command or no contest context" }
@@ -398,7 +418,7 @@ function M.handle_command(opts)
 				)
 			end
 
-			setup_problem(cmd.contest, cmd.problem)
+			setup_problem(cmd.contest, cmd.problem, cmd.language)
 		end
 		return
 	end
@@ -412,16 +432,16 @@ function M.handle_command(opts)
 					vim.log.levels.WARN
 				)
 			end
-			setup_problem(cmd.problem)
+			setup_problem(cmd.problem, nil, cmd.language)
 		end
 		return
 	end
 
 	if cmd.type == "problem_switch" then
 		if state.platform == "cses" then
-			setup_problem(cmd.problem)
+			setup_problem(cmd.problem, nil, cmd.language)
 		else
-			setup_problem(state.contest_id, cmd.problem)
+			setup_problem(state.contest_id, cmd.problem, cmd.language)
 		end
 		return
 	end

lua/cp/problem.lua

@@ -15,10 +15,11 @@ local M = {}
 ---@param contest_id string
 ---@param problem_id? string
 ---@param config cp.Config
+---@param language? string
 ---@return ProblemContext
-function M.create_context(contest, contest_id, problem_id, config)
+function M.create_context(contest, contest_id, problem_id, config, language)
 	local filename_fn = config.filename or require("cp.config").default_filename
-	local source_file = filename_fn(contest, contest_id, problem_id)
+	local source_file = filename_fn(contest, contest_id, problem_id, config, language)
 	local base_name = vim.fn.fnamemodify(source_file, ":t:r")
 
 	return {

lua/cp/snippets.lua

@@ -10,102 +10,113 @@ function M.setup(config)
 
 	local s, i, fmt = ls.snippet, ls.insert_node, require("luasnip.extras.fmt").fmt
 
-	local default_snippets = {
+	local filetype_to_language = {
-		s(
+		cc = "cpp",
-			"codeforces",
+		c = "cpp",
-			fmt(
+		py = "python",
-				[[#include <bits/stdc++.h>
+		py3 = "python",
+	}
+
+	local language_to_filetype = {}
+	for ext, lang in pairs(filetype_to_language) do
+		language_to_filetype[lang] = ext
+	end
+
+	local template_definitions = {
+		cpp = {
+			codeforces = [[#include <bits/stdc++.h>
 
 using namespace std;
 
-void solve() {{
+void solve() {
   {}
-}}
+}
 
-int main() {{
+int main() {
   std::cin.tie(nullptr)->sync_with_stdio(false);
 
   int tc = 1;
   std::cin >> tc;
 
-  for (int t = 0; t < tc; ++t) {{
+  for (int t = 0; t < tc; ++t) {
     solve();
-  }}
+  }
 
   return 0;
-}}]],
+}]],
-				{ i(1) }
-			)
-		),
 
-		s(
+			atcoder = [[#include <bits/stdc++.h>
-			"atcoder",
-			fmt(
-				[[#include <bits/stdc++.h>
 
 using namespace std;
 
-void solve() {{
+void solve() {
   {}
-}}
+}
 
-int main() {{
+int main() {
   std::cin.tie(nullptr)->sync_with_stdio(false);
 
 #ifdef LOCAL
   int tc;
   std::cin >> tc;
 
-  for (int t = 0; t < tc; ++t) {{
+  for (int t = 0; t < tc; ++t) {
     solve();
-  }}
+  }
 #else
   solve();
 #endif
 
   return 0;
-}}]],
+}]],
-				{ i(1) }
-			)
-		),
 
-		s(
+			cses = [[#include <bits/stdc++.h>
-			"cses",
-			fmt(
-				[[#include <bits/stdc++.h>
 
 using namespace std;
 
-int main() {{
+int main() {
   std::cin.tie(nullptr)->sync_with_stdio(false);
 
   {}
 
   return 0;
-}}]],
+}]],
-				{ i(1) }
+		},
-			)
+
-		),
+		python = {
+			codeforces = [[def solve():
+    {}
+
+if __name__ == "__main__":
+    tc = int(input())
+    for _ in range(tc):
+        solve()]],
+
+			atcoder = [[def solve():
+    {}
+
+if __name__ == "__main__":
+    solve()]],
+
+			cses = [[{}]],
+		},
 	}
 
-	local default_map = {}
+	for language, filetype in pairs(language_to_filetype) do
-	for _, snippet in pairs(default_snippets) do
+		local snippets = {}
-		default_map[snippet.trigger] = snippet
+
+		for contest, template in pairs(template_definitions[language] or {}) do
+			table.insert(snippets, s(contest, fmt(template, { i(1) })))
+		end
+
+		for _, snippet in ipairs(config.snippets or {}) do
+			if snippet.filetype == filetype then
+				table.insert(snippets, snippet)
+			end
+		end
+
+		ls.add_snippets(filetype, snippets)
 	end
-
-	local user_map = {}
-	for _, snippet in pairs(config.snippets or {}) do
-		user_map[snippet.trigger] = snippet
-	end
-
-	local merged_map = vim.tbl_extend("force", default_map, user_map)
-
-	local all_snippets = {}
-	for _, snippet in pairs(merged_map) do
-		table.insert(all_snippets, snippet)
-	end
-
-	ls.add_snippets("cpp", all_snippets)
 end
 
 return M

plugin/cp.lua

@@ -11,15 +11,41 @@ vim.api.nvim_create_user_command("CP", function(opts)
 	cp.handle_command(opts)
 end, {
 	nargs = "*",
+	desc = "Competitive programming helper",
 	complete = function(ArgLead, CmdLine, _)
+		local filetype_to_language = {
+			cc = "cpp",
+			c = "cpp",
+			py = "python",
+			py3 = "python",
+		}
+
+		local languages = vim.tbl_keys(vim.tbl_add_reverse_lookup(filetype_to_language))
+
+		if ArgLead:match("^--lang=") then
+			local lang_completions = {}
+			for _, lang in ipairs(languages) do
+				table.insert(lang_completions, "--lang=" .. lang)
+			end
+			return vim.tbl_filter(function(completion)
+				return completion:find(ArgLead, 1, true) == 1
+			end, lang_completions)
+		end
+
+		if ArgLead == "--lang" then
+			return { "--lang" }
+		end
+
 		local args = vim.split(vim.trim(CmdLine), "%s+")
 		local num_args = #args
 		if CmdLine:sub(-1) == " " then
 			num_args = num_args + 1
 		end
 
+		local lang_flag_present = vim.tbl_contains(args, "--lang") or vim.iter(args):any(function(arg) return arg:match("^--lang=") end)
+
 		if num_args == 2 then
-			local candidates = {}
+			local candidates = { "--lang" }
 			vim.list_extend(candidates, platforms)
 			vim.list_extend(candidates, actions)
 			local cp = require("cp")
@@ -37,13 +63,17 @@ end, {
 			return vim.tbl_filter(function(cmd)
 				return cmd:find(ArgLead, 1, true) == 1
 			end, candidates)
-		elseif num_args == 4 then
+		elseif args[#args-1] == "--lang" then
+			return vim.tbl_filter(function(lang)
+				return lang:find(ArgLead, 1, true) == 1
+			end, languages)
+		elseif num_args == 4 and not lang_flag_present 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 = {}
+					local candidates = { "--lang" }
 					for _, problem in ipairs(contest_data.problems) do
 						table.insert(candidates, problem.id)
 					end