2525
2626type_errors = []
2727
28- def parse_type (annotation ):
28+ class ClassInfo :
29+ """存储自定义类的信息"""
30+ def __init__ (self , name :str ):
31+ self .name = name
32+ self .fields : dict [str , ast .AST ] = {} # 字段名 -> 类型注解
33+ self .bases : list [str ] = [] # 基类名
34+
35+ def add_field (self , field_name : str , annotation : ast .AST ):
36+ """添加字段"""
37+ self .fields [field_name ] = annotation
38+
39+ def add_base (self , base_name : str ):
40+ """添加基类"""
41+ self .bases .append (base_name )
42+
43+
44+ def parse_type (annotation , class_map ):
2945 """解析参数类型"""
3046 global type_errors
3147
@@ -34,79 +50,133 @@ def parse_type(annotation):
3450 return "invalid" , None , True
3551
3652 elif isinstance (annotation , ast .Name ):
37- if annotation .id in TYPE_MAP :
53+ if annotation .id in TYPE_MAP : # 基础类型
3854 return TYPE_MAP [annotation .id ], None , True
39- else :
55+ elif annotation .id in class_map : # 自定义类型
56+ class_info = class_map [annotation .id ]
57+ properties = {}
58+ for field_name , field_annotation in class_info .fields .items ():
59+ field_type , field_items , _ = parse_type (field_annotation , class_map )
60+
61+ if field_type == "invalid" :
62+ type_errors .append (f"类 { annotation .id } 的字段 { field_name } 类型无效" )
63+ return "invalid" , None , True
64+
65+ field_schema = {"type" : field_type }
66+ if field_type == "array" :
67+ field_schema ["items" ] = field_items if field_items else {}
68+ elif field_type == "object" and field_items :
69+ if "properties" in field_items :
70+ field_schema ["properties" ] = field_items ["properties" ]
71+ if "anyOf" in field_items :
72+ field_schema ["anyOf" ] = field_items ["anyOf" ]
73+ if "additionalProperties" in field_items :
74+ field_schema ["additionalProperties" ] = field_items ["additionalProperties" ]
75+
76+ properties [field_name ] = field_schema
77+ return "object" , {"type" :"object" , "properties" :properties }, True
78+ else : # 未知类型
4079 type_errors .append (f"不支持的类型: { annotation .id } " )
4180 return "invalid" , None , True
4281
43- elif isinstance (annotation , ast .Constant ) and annotation .value is None :
82+ elif isinstance (annotation , ast .Constant ) and annotation .value is None : # None
4483 return "null" , None , False
4584
46- elif isinstance (annotation , ast .Subscript ):
47- if isinstance (annotation .value , ast .Name ):
48- container = annotation .value .id
85+ elif isinstance (annotation , ast .Subscript ) and isinstance (annotation .value , ast .Name ): # 容器类型
86+ container = annotation .value .id
4987
50- # List[int]
51- if container in ("list" , "List" ):
52- item_type , _ , _ = parse_type (annotation .slice )
53- if item_type == "invalid" :
54- type_errors .append (f"不支持的列表元素类型: { annotation .slice } " )
88+ # List[int]
89+ if container in ("list" , "List" ):
90+ item_type , item_schema , _ = parse_type (annotation .slice , class_map )
91+ if item_type == "invalid" :
92+ type_errors .append (f"不支持的列表元素类型: { annotation .slice } " )
93+ return "invalid" , None , True
94+ items = item_schema if item_schema else {"type" : item_type }
95+ return "array" , items , True
96+
97+ # Dict[str, int] → object
98+ elif container in ("dict" , "Dict" ):
99+ if isinstance (annotation .slice , ast .Tuple ) and len (annotation .slice .elts ) == 2 :
100+ key_annot , value_annot = annotation .slice .elts
101+ key_type , _ , _ = parse_type (key_annot , class_map )
102+ if key_type != "string" :
103+ type_errors .append (f"Dict 的键类型必须是 string,实际是 { key_type } " )
55104 return "invalid" , None , True
56- return "array" , {"type" : item_type }, True
105+ value_type , value_schema , _ = parse_type (value_annot , class_map )
106+ items = value_schema if value_schema else {"type" : value_type }
107+ return "object" , {"additionalProperties" : items }, True
57108
58- # Dict[str, int] → object
59- elif container in ( "dict" , "Dict" ) :
60- return "object" , None , True
109+ # Optional[ int]
110+ elif container == "Optional" :
111+ inner_type , inner_items , _ = parse_type ( annotation . slice , class_map )
61112
62- # Optional[int]
63- elif container == "Optional" :
64- inner_type , inner_items , _ = parse_type (annotation .slice )
113+ if inner_type == "invalid" :
114+ type_errors .append (f"不支持的Optional类型: { annotation .slice } " )
115+ return "invalid" , None , False
116+ return inner_type , inner_items , False
117+
118+ # Union[str, int]
119+ elif container == "Union" :
120+ if isinstance (annotation .slice , ast .Tuple ):
121+ schemas = []
122+ for elt in annotation .slice .elts :
123+ elt_type , elt_items , _ = parse_type (elt , class_map )
124+ if elt_type == "invalid" :
125+ type_errors .append (f"不支持的 Union 元素类型: { ast .dump (elt )} " )
126+ return "invalid" , None , True
127+ schema = {"type" : elt_type }
128+ if elt_items :
129+ schema .update (elt_items )
130+ schemas .append (schema )
131+ return "object" , {"anyOf" : schemas }, True
132+ else :
133+ inner_type , inner_items , _ = parse_type (annotation .slice , class_map )
65134 if inner_type == "invalid" :
66- type_errors .append (f"不支持的Optional类型: { annotation .slice } " )
67- return "invalid" , None , False
68- return inner_type , inner_items , False
69-
70- # Union[str, int]
71- elif container == "Union" :
72- return "object" , None , True
73-
74- # Tuple[str]
75- elif container in ("tuple" , "Tuple" ):
76- items = []
77- if isinstance (annotation .slice , ast .Tuple ):
78- for elt in annotation .slice .elts :
79- item_type , _ , _ = parse_type (elt )
80- if item_type == "invalid" :
81- type_errors .append (f"不支持的元组元素类型: { ast .dump (elt )} " )
82- return "invalid" , None , True
83- items .append ({"type" :item_type })
84- return "array" , f"{ items } " , True
85- else :
86- item_type , _ , _ = parse_type (annotation .slice )
135+ type_errors .append (f"不支持的 Union 类型: { ast .dump (annotation .slice )} " )
136+ return "invalid" , None , True
137+ schema = {"type" : inner_type }
138+ if inner_items :
139+ schema .update (inner_items )
140+ return "object" , {"anyOf" : [schema ]}, True
141+
142+ # Tuple[str]
143+ elif container in ("tuple" , "Tuple" ):
144+ if isinstance (annotation .slice , ast .Tuple ):
145+ tuple_items = []
146+ for elt in annotation .slice .elts :
147+ item_type , item_schema , _ = parse_type (elt , class_map )
87148 if item_type == "invalid" :
88- type_errors .append (f"不支持的元组元素类型: { ast .dump (annotation . slice )} " )
149+ type_errors .append (f"不支持的元组元素类型: { ast .dump (elt )} " )
89150 return "invalid" , None , True
90- return "array" , {"type" :item_type }, True
91-
92- # Set[int]
93- elif container in ("set" , "Set" ):
94- item_type , _ , _ = parse_type (annotation .slice )
151+ tuple_items .append (item_schema if item_schema else {"type" : item_type })
152+ # 返回固定长度 tuple 的 items 列表
153+ return "array" , {"items" : tuple_items }, True
154+ else :
155+ # 单元素 Tuple
156+ item_type , item_schema , _ = parse_type (annotation .slice , class_map )
95157 if item_type == "invalid" :
96- type_errors .append (f"不支持的集合元素类型 : { annotation .slice } " )
158+ type_errors .append (f"不支持的元组元素类型 : { ast . dump ( annotation .slice ) } " )
97159 return "invalid" , None , True
98- return "array" , {"type" : item_type }, True
99-
100-
101- else :
102- type_errors .append (f"不支持的容器类型: { container } " )
160+ return "array" , {"items" : item_schema if item_schema else {"type" : item_type }}, True
161+
162+ # Set[int]
163+ elif container in ("set" , "Set" ):
164+ item_type , item_schema , _ = parse_type (annotation .slice , class_map )
165+ if item_type == "invalid" :
166+ type_errors .append (f"不支持的集合元素类型: { annotation .slice } " )
103167 return "invalid" , None , True
168+ items = item_schema if item_schema else {"type" : item_type }
169+ return "array" , items , True
170+
171+ else :
172+ type_errors .append (f"不支持的容器类型: { container } " )
173+ return "invalid" , None , True
104174
105175 type_errors .append (f"无法识别的类型: { ast .dump (annotation )} " )
106176 return "invalid" , None , True
107177
108178
109- def parse_parameters (args ):
179+ def parse_parameters (args , class_map ):
110180 """解析函数参数"""
111181 properties = {}
112182 order = []
@@ -115,35 +185,81 @@ def parse_parameters(args):
115185 for arg in args .args :
116186 arg_name = arg .arg
117187 order .append (arg_name )
118- arg_type , items , is_required = parse_type (arg .annotation )
188+ arg_type , items , is_required = parse_type (arg .annotation , class_map )
119189 # 定义参数
120190 prop_def = {
121191 "defaultValue" : "" ,
122192 "description" : f"参数 { arg_name } " ,
123193 "name" : arg_name ,
124194 "type" : arg_type ,
125- ** ({"items" : items } if items else {}),
126- "examples" : "" ,
127- "required" : is_required ,
128195 }
196+ if arg_type == "array" and items :
197+ if "items" in items :
198+ prop_def ["items" ] = items ["items" ]
199+ else :
200+ arr_items = {"type" : items .get ("type" , "object" )}
201+ if "properties" in items :
202+ arr_items ["properties" ] = items ["properties" ]
203+ if "anyOf" in items :
204+ arr_items ["anyOf" ] = items ["anyOf" ]
205+ if "additionalProperties" in items :
206+ arr_items ["additionalProperties" ] = items ["additionalProperties" ]
207+ prop_def ["items" ] = arr_items
208+
209+ if arg_type == "object" and items :
210+ if "properties" in items :
211+ prop_def ["properties" ] = items ["properties" ]
212+ if "anyOf" in items :
213+ prop_def ["anyOf" ] = items ["anyOf" ]
214+ if "additionalProperties" in items :
215+ prop_def ["additionalProperties" ] = items ["additionalProperties" ]
216+
217+ prop_def ["examples" ] = ""
218+ prop_def ["required" ] = is_required
219+
129220 properties [arg_name ] = prop_def
130221 if is_required :
131222 required .append (arg_name )
132223 return properties , order , required
133224
134225
135- def parse_return (annotation ):
226+ def parse_return (annotation , custom_classes ):
136227 """解析返回值类型"""
137228 if not annotation :
138229 return {"type" : "string" , "convertor" : "" }
139230
140- return_type , items , _ = parse_type (annotation )
141- ret = {
142- "type" : return_type ,
143- ** ({"items" : items } if items else {}),
144- "convertor" : ""
145- }
146- return ret
231+ return_type , items , _ = parse_type (annotation , custom_classes )
232+ if return_type == "array" :
233+ ret = {"type" : "array" }
234+ if items :
235+ if "items" in items :
236+ ret ["items" ] = items ["items" ]
237+ else :
238+ arr_items = {"type" : items .get ("type" , "object" )}
239+ if isinstance (items , dict ) and "properties" in items :
240+ arr_items ["properties" ] = items ["properties" ]
241+ if isinstance (items , dict ) and "anyOf" in items :
242+ arr_items ["anyOf" ] = items ["anyOf" ]
243+ if isinstance (items , dict ) and "additionalProperties" in items :
244+ arr_items ["additionalProperties" ] = items ["additionalProperties" ]
245+ ret ["items" ] = arr_items
246+ ret ["convertor" ] = ""
247+ return ret
248+
249+ elif return_type == "object" :
250+ ret = {"type" : "object" }
251+ if items and isinstance (items , dict ):
252+ if "properties" in items :
253+ ret ["properties" ] = items ["properties" ]
254+ if "anyOf" in items :
255+ ret ["anyOf" ] = items ["anyOf" ]
256+ if "additionalProperties" in items :
257+ ret ["additionalProperties" ] = items ["additionalProperties" ]
258+ ret ["convertor" ] = ""
259+ return ret
260+
261+ else :
262+ return {"type" : return_type , "convertor" : "" }
147263
148264
149265def parse_python_file (file_path : Path ):
@@ -155,11 +271,26 @@ def parse_python_file(file_path: Path):
155271 py_name = file_path .stem
156272 definitions = []
157273 tool_groups = []
158-
274+ class_map : dict [str , ClassInfo ] = {} # 类名, 类信息
275+ # 收集自定义类
276+ for node in tree .body :
277+ if isinstance (node , ast .ClassDef ):
278+ class_info = ClassInfo (node .name )
279+ for base in node .bases :
280+ if isinstance (base , ast .Name ):
281+ class_info .add_base (base .id )
282+ for subnode in node .body :
283+ if isinstance (subnode , ast .FunctionDef ) and subnode .name == "__init__" :
284+ for arg in subnode .args .args [1 :]: # 跳过 self
285+ arg_name = arg .arg
286+ class_info .add_field (arg_name , arg .annotation )
287+ class_map [node .name ] = class_info
288+ # 解析函数定义
159289 for node in tree .body :
160290 if isinstance (node , ast .FunctionDef ):
161291 func_name = node .name
162- # 默认描述
292+
293+ # 获取描述
163294 description = f"执行 { func_name } 方法"
164295 if node .body and isinstance (node .body [0 ], ast .Expr ):
165296 expr_value = node .body [0 ].value
@@ -185,8 +316,8 @@ def parse_python_file(file_path: Path):
185316 continue
186317
187318 # 解析参数和返回值
188- properties , order , required = parse_parameters (node .args )
189- return_schema = parse_return (node .returns )
319+ properties , order , required = parse_parameters (node .args , class_map )
320+ return_schema = parse_return (node .returns , class_map )
190321
191322 # definition schema
192323 definition_schema = {
@@ -213,6 +344,9 @@ def parse_python_file(file_path: Path):
213344 "name" : v ["name" ],
214345 "type" : v ["type" ],
215346 ** ({"items" : v ["items" ]} if "items" in v else {}),
347+ ** ({"properties" : v ["properties" ]} if "properties" in v else {}),
348+ ** ({"anyOf" : v ["anyOf" ]} if "anyOf" in v else {}),
349+ ** ({"additionalProperties" : v ["additionalProperties" ]} if "additionalProperties" in v else {}),
216350 "required" : False , # 工具里参数默认非必填
217351 }
218352 for k , v in properties .items ()
@@ -225,6 +359,9 @@ def parse_python_file(file_path: Path):
225359 "description" : f"{ func_name } 函数的返回值" ,
226360 "type" : return_schema ["type" ],
227361 ** ({"items" : return_schema ["items" ]} if "items" in return_schema else {}),
362+ ** ({"properties" : return_schema ["properties" ]} if "properties" in return_schema else {}),
363+ ** ({"anyOf" : return_schema ["anyOf" ]} if "anyOf" in return_schema else {}),
364+ ** ({"additionalProperties" : return_schema ["additionalProperties" ]} if "additionalProperties" in return_schema else {}),
228365 "convertor" : "" ,
229366 "examples" : "" ,
230367 },
0 commit comments