1212from ..utils .compiler import run_binary , run_binary_with_args
1313from ..utils .platform import get_exe_extension
1414from .base import Tool , ToolResult
15- from .mixins import BuildToolMixin
15+ from .mixins import BuildToolMixin , resolve_source
1616
1717
1818class GeneratorBuildTool (Tool , BuildToolMixin ):
@@ -74,55 +74,43 @@ async def execute(
7474 compiler : str = "g++" ,
7575 ) -> ToolResult :
7676 """执行 Generator 构建。"""
77- # 解析源代码:source_path 优先于 code
78- source_dir = None
79- if source_path :
80- if not os .path .isabs (source_path ):
81- source_path = os .path .join (problem_dir , source_path )
82- if not os .path .exists (source_path ):
83- return ToolResult .fail (f"Source file not found: { source_path } " )
84- try :
85- with open (source_path , encoding = "utf-8" ) as f :
86- code = f .read ()
87- except UnicodeDecodeError :
88- try :
89- with open (source_path , encoding = "latin-1" ) as f :
90- code = f .read ()
91- except Exception as e :
92- return ToolResult .fail (f"Failed to read source file: { e } " )
93- source_dir = os .path .dirname (os .path .abspath (source_path ))
94- elif code is None :
95- return ToolResult .fail ("Either 'code' or 'source_path' must be provided" )
77+ resolved , err = resolve_source (problem_dir , code , source_path )
78+ if resolved is None :
79+ return err
9680
9781 os .makedirs (problem_dir , exist_ok = True )
98-
99- # 保存到 files/ 子目录
10082 files_dir = os .path .join (problem_dir , "files" )
10183 os .makedirs (files_dir , exist_ok = True )
10284
103- source_path = os .path .join (files_dir , "gen.cpp" )
85+ canonical_path = os .path .join (files_dir , "gen.cpp" )
10486 try :
105- with open (source_path , "w" , encoding = "utf-8" ) as f :
106- f .write (code )
87+ with open (canonical_path , "w" , encoding = "utf-8" ) as f :
88+ f .write (resolved . code )
10789 except Exception as e :
10890 return ToolResult .fail (f"Failed to save code: { str (e )} " )
10991
11092 exe_ext = get_exe_extension ()
11193 binary_path = os .path .join (files_dir , f"gen{ exe_ext } " )
11294
113- include_dirs = [source_dir ] if source_dir else None
114- compile_result = await self .build (source_path , binary_path , compiler = compiler , include_dirs = include_dirs )
95+ compile_source = resolved .original_source_path or canonical_path
96+ include_dirs = [resolved .include_dir ] if resolved .include_dir else None
97+ compile_result = await self .build (compile_source , binary_path , compiler = compiler , include_dirs = include_dirs )
11598
11699 if not compile_result .success :
117100 return ToolResult .fail (
118101 f"Compilation failed: { compile_result .error } " ,
119- source_path = source_path ,
102+ source_path = compile_source ,
103+ canonical_path = canonical_path ,
120104 compile_log = compile_result .stderr ,
121105 )
122106
107+ binary_size = os .path .getsize (binary_path ) if os .path .exists (binary_path ) else 0
108+
123109 return ToolResult .ok (
124- source_path = source_path ,
110+ source_path = compile_source ,
111+ canonical_path = canonical_path ,
125112 binary_path = binary_path ,
113+ binary_size = binary_size ,
126114 compile_log = compile_result .stderr ,
127115 message = "Generator built successfully" ,
128116 )
@@ -207,6 +195,12 @@ def input_schema(self) -> dict:
207195 "description" : "T 最大值" ,
208196 "default" : 1 ,
209197 },
198+ "extra_args" : {
199+ "type" : "array" ,
200+ "items" : {"type" : "string" },
201+ "description" : "附加命令行参数,追加在标准 6 参数之后传递给 generator" ,
202+ "default" : [],
203+ },
210204 },
211205 "required" : ["problem_dir" , "strategies" ],
212206 }
@@ -222,9 +216,11 @@ async def execute(
222216 n_max : int = 100000 ,
223217 t_min : int = 1 ,
224218 t_max : int = 1 ,
219+ extra_args : list [str ] | None = None ,
225220 ) -> ToolResult :
226221 """执行数据生成。"""
227222 exe_ext = get_exe_extension ()
223+ extra_args = extra_args or []
228224
229225 # 检查 generator - 优先查找 files/ 子目录
230226 gen_exe = os .path .join (problem_dir , "files" , f"gen{ exe_ext } " )
@@ -262,8 +258,8 @@ async def execute(
262258 type_param = strategy_type_map .get (strategy , "2" )
263259
264260 # 运行 generator
265- # gen.exe <seed> <type> <n_min> <n_max> <t_min> <t_max>
266- cmd_args = [str (seed ), type_param , str (n_min ), str (n_max ), str (t_min ), str (t_max )]
261+ # gen.exe <seed> <type> <n_min> <n_max> <t_min> <t_max> [extra_args...]
262+ cmd_args = [str (seed ), type_param , str (n_min ), str (n_max ), str (t_min ), str (t_max )] + extra_args
267263
268264 gen_result = await run_binary_with_args (
269265 gen_exe ,
0 commit comments