Skip to content

Commit ebc9066

Browse files
authored
[fit] build命令支持自定义结构体参数 (#323)
* [fit] build命令支持自定义结构体 * [fit] 完善build 逻辑
1 parent 6da6f2c commit ebc9066

File tree

1 file changed

+206
-69
lines changed
  • framework/fit/python/fit_cli/utils

1 file changed

+206
-69
lines changed

framework/fit/python/fit_cli/utils/build.py

Lines changed: 206 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,23 @@
2525

2626
type_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

149265
def 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

Comments
 (0)