feat(tools): 优化 JSON 示例生成函数

- 增加 include_optional 参数,决定是否包含可选字段
- 添加 list_length 参数,用于控制列表字段的示例长度
- 在列表示例中添加省略标记,更直观展示多元素列表
- 优化字典字段的示例生成逻辑
This commit is contained in:
carry 2025-04-19 21:07:00 +08:00
parent 9fc3ab904b
commit 1e829c9268

View File

@ -3,17 +3,19 @@ from typing import Any, Dict, List, Optional, Union, get_args, get_origin
import json
from datetime import datetime, date
def generate_example_json(model: type[BaseModel]) -> str:
def generate_example_json(model: type[BaseModel], include_optional: bool = False,list_length = 2) -> str:
"""
根据 Pydantic V2 模型生成示例 JSON 数据结构
"""
def _generate_example(field_type: Any) -> Any:
origin = get_origin(field_type)
args = get_args(field_type)
if origin is list or origin is List:
return [_generate_example(args[0])] if args else []
# 生成多个元素,这里生成 3 个
result = [_generate_example(args[0]) for _ in range(list_length)] if args else []
result.append("......")
return result
elif origin is dict or origin is Dict:
if len(args) == 2:
return {"key": _generate_example(args[1])}
@ -35,21 +37,22 @@ def generate_example_json(model: type[BaseModel]) -> str:
elif field_type is date:
return date.today().isoformat()
elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
return json.loads(generate_example_json(field_type))
return json.loads(generate_example_json(field_type, include_optional))
else:
# 处理直接类型注解(非泛型)
if field_type is type(None):
return None
try:
if issubclass(field_type, BaseModel):
return json.loads(generate_example_json(field_type))
return json.loads(generate_example_json(field_type, include_optional))
except TypeError:
pass
return "unknown"
example_data = {}
for field_name, field in model.model_fields.items():
example_data[field_name] = _generate_example(field.annotation)
if include_optional or not isinstance(field.default, type(None)):
example_data[field_name] = _generate_example(field.annotation)
return json.dumps(example_data, indent=2, default=str)