354 lines
12 KiB
Python
354 lines
12 KiB
Python
|
"""
|
||
|
|
模型生成脚本 — 从 config/workers/<plugin_name>/*.toml 生成 Pydantic 模型
|
||
|
|
|
||
|
|
用法:
|
||
|
|
uv run gen_models config/workers/example src/plugins/example
|
||
|
|
|
||
|
|
工作流:
|
||
|
|
1. 编写 config/workers/<plugin_name>/*.toml 配置文件
|
||
|
|
2. 执行 gen_models 生成 Pydantic 模型
|
||
|
|
3. 在业务代码中 import 生成的模型
|
||
|
|
|
||
|
|
功能:
|
||
|
|
1. 扫描配置目录下的所有 .toml 文件
|
||
|
|
2. 解析每个路由的 params 和 response_model
|
||
|
|
3. 生成 Pydantic 请求/响应模型
|
||
|
|
4. 写入 src/plugins/<plugin_name>/models/__init__.py 和 models/pydantic.py
|
||
|
|
"""
|
||
|
|
|
||
|
|
import sys
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
def parse_type_str(type_str: str) -> str:
|
||
|
|
"""将 TOML 中的类型字符串映射为 Python 类型注解"""
|
||
|
|
mapping = {
|
||
|
|
"float": "float",
|
||
|
|
"int": "int",
|
||
|
|
"str": "str",
|
||
|
|
"bool": "bool",
|
||
|
|
"string": "str",
|
||
|
|
"integer": "int",
|
||
|
|
"number": "float",
|
||
|
|
"boolean": "bool",
|
||
|
|
"any": "Any",
|
||
|
|
}
|
||
|
|
|
||
|
|
if type_str.startswith("list[") and type_str.endswith("]"):
|
||
|
|
inner = type_str[5:-1]
|
||
|
|
inner_py = mapping.get(inner, inner)
|
||
|
|
return f"list[{inner_py}]"
|
||
|
|
|
||
|
|
if type_str == "dict":
|
||
|
|
return "dict[str, Any]"
|
||
|
|
|
||
|
|
if type_str.startswith("dict[") and type_str.endswith("]"):
|
||
|
|
inner = type_str[5:-1]
|
||
|
|
if "," in inner:
|
||
|
|
parts = [p.strip() for p in inner.split(",", 1)]
|
||
|
|
inner_py = mapping.get(parts[-1], parts[-1])
|
||
|
|
else:
|
||
|
|
inner_py = mapping.get(inner, inner)
|
||
|
|
return f"dict[str, {inner_py}]"
|
||
|
|
|
||
|
|
if type_str.startswith("Optional[") and type_str.endswith("]"):
|
||
|
|
inner = type_str[9:-1]
|
||
|
|
inner_py = mapping.get(inner, inner)
|
||
|
|
return f"Optional[{inner_py}]"
|
||
|
|
|
||
|
|
return mapping.get(type_str, type_str)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_params(params: list[str]) -> list[dict[str, str]]:
|
||
|
|
"""解析 params 列表，如 ['a: float', 'b: float']"""
|
||
|
|
result = []
|
||
|
|
for p in params:
|
||
|
|
if ":" in p:
|
||
|
|
name, type_str = p.split(":", 1)
|
||
|
|
result.append({"name": name.strip(), "type": parse_type_str(type_str.strip())})
|
||
|
|
else:
|
||
|
|
result.append({"name": p.strip(), "type": "str"})
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def parse_response_model(response_model: dict[str, Any]) -> list[dict[str, Any]]:
|
||
|
|
"""解析 response_model 字典，返回字段列表"""
|
||
|
|
fields: list[dict[str, Any]] = []
|
||
|
|
for field_name, field_spec in response_model.items():
|
||
|
|
entry: dict[str, Any] = {"name": field_name}
|
||
|
|
if isinstance(field_spec, dict):
|
||
|
|
entry["type"] = parse_type_str(field_spec.get("type", "str"))
|
||
|
|
entry["required"] = field_spec.get("required", True)
|
||
|
|
if "default" in field_spec:
|
||
|
|
entry["default"] = field_spec["default"]
|
||
|
|
else:
|
||
|
|
entry["type"] = parse_type_str(str(field_spec))
|
||
|
|
entry["required"] = True
|
||
|
|
fields.append(entry)
|
||
|
|
return fields
|
||
|
|
|
||
|
|
|
||
|
|
def _model_name(base: str, method: str, suffix: str) -> str:
|
||
|
|
"""生成模型名: e.g. add + GET + Request -> AddGetRequest, add + MQTT + Request -> AddMqttRequest"""
|
||
|
|
method_title = "Mqtt" if method == "MQTT" else method.title()
|
||
|
|
return f"{base.title().replace('-', '').replace('_', '')}{method_title}{suffix}"
|
||
|
|
|
||
|
|
|
||
|
|
def generate_pydantic_code(
|
||
|
|
plugin_name: str,
|
||
|
|
routes: list[dict[str, Any]],
|
||
|
|
) -> str:
|
||
|
|
"""生成 pydantic.py 文件内容"""
|
||
|
|
# 收集全部字段类型，按需生成 typing 导入（避免未使用导入触发 lint）
|
||
|
|
type_strings: list[str] = []
|
||
|
|
has_optional_field = False
|
||
|
|
for route in routes:
|
||
|
|
for p in route.get("parsed_params", []):
|
||
|
|
type_strings.append(p.get("type", "str"))
|
||
|
|
for f in route.get("parsed_response", []):
|
||
|
|
type_strings.append(f.get("type", "str"))
|
||
|
|
# 非必填字段渲染为 Optional[...]
|
||
|
|
if not f.get("required", True):
|
||
|
|
has_optional_field = True
|
||
|
|
typing_imports = []
|
||
|
|
if has_optional_field or any("Optional" in t for t in type_strings):
|
||
|
|
typing_imports.append("Optional")
|
||
|
|
if any("Any" in t for t in type_strings):
|
||
|
|
typing_imports.append("Any")
|
||
|
|
|
||
|
|
lines = []
|
||
|
|
lines.append('"""')
|
||
|
|
lines.append(f"自动生成的 Pydantic 模型 — {plugin_name}")
|
||
|
|
lines.append(f"来源: config/workers/{plugin_name}/")
|
||
|
|
lines.append('"""')
|
||
|
|
lines.append("")
|
||
|
|
lines.append("from __future__ import annotations")
|
||
|
|
lines.append("")
|
||
|
|
if typing_imports:
|
||
|
|
lines.append(f"from typing import {', '.join(typing_imports)}")
|
||
|
|
lines.append("from pydantic import BaseModel, Field")
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
seen_models: set[str] = set()
|
||
|
|
|
||
|
|
for route in routes:
|
||
|
|
route_name = route["name"]
|
||
|
|
method = route["method"]
|
||
|
|
summary = route.get("summary", "")
|
||
|
|
|
||
|
|
# --- Request Model ---
|
||
|
|
params = route.get("parsed_params", [])
|
||
|
|
if params:
|
||
|
|
model_name = _model_name(route_name, method, "Request")
|
||
|
|
if model_name not in seen_models:
|
||
|
|
seen_models.add(model_name)
|
||
|
|
lines.append(f"class {model_name}(BaseModel):")
|
||
|
|
lines.append(f' """{summary} — 请求体"""')
|
||
|
|
for p in params:
|
||
|
|
lines.append(f" {p['name']}: {p['type']} = Field(...)")
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
# --- Response Model ---
|
||
|
|
response_fields = route.get("parsed_response", [])
|
||
|
|
if response_fields:
|
||
|
|
model_name = _model_name(route_name, method, "Response")
|
||
|
|
if model_name not in seen_models:
|
||
|
|
seen_models.add(model_name)
|
||
|
|
lines.append(f"class {model_name}(BaseModel):")
|
||
|
|
lines.append(f' """{summary} — 响应体"""')
|
||
|
|
for f in response_fields:
|
||
|
|
field_type = f["type"]
|
||
|
|
if f.get("required", True):
|
||
|
|
if "default" in f:
|
||
|
|
default_repr = (
|
||
|
|
json.dumps(f["default"])
|
||
|
|
if not isinstance(f["default"], str)
|
||
|
|
else f"'{f['default']}'"
|
||
|
|
)
|
||
|
|
lines.append(f" {f['name']}: {field_type} = Field(default={default_repr})")
|
||
|
|
else:
|
||
|
|
lines.append(f" {f['name']}: {field_type} = Field(...)")
|
||
|
|
else:
|
||
|
|
if "default" in f:
|
||
|
|
default_repr = (
|
||
|
|
json.dumps(f["default"])
|
||
|
|
if not isinstance(f["default"], str)
|
||
|
|
else f"'{f['default']}'"
|
||
|
|
)
|
||
|
|
lines.append(f" {f['name']}: Optional[{field_type}] = Field(default={default_repr})")
|
||
|
|
else:
|
||
|
|
lines.append(f" {f['name']}: Optional[{field_type}] = Field(default=None)")
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
# --- Union type for all responses ---
|
||
|
|
response_models = [m for m in seen_models if m.endswith("Response")]
|
||
|
|
if response_models:
|
||
|
|
plugin_title = plugin_name.title().replace("-", "").replace("_", "")
|
||
|
|
lines.append("")
|
||
|
|
lines.append("# 所有响应模型的联合类型")
|
||
|
|
lines.append(f"{plugin_title}Response = " + " | ".join(response_models))
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
|
||
|
|
def generate_init_code(
|
||
|
|
plugin_name: str,
|
||
|
|
routes: list[dict[str, Any]],
|
||
|
|
) -> str:
|
||
|
|
"""生成 __init__.py 文件内容"""
|
||
|
|
lines = []
|
||
|
|
lines.append('"""')
|
||
|
|
lines.append(f"模型导出 — {plugin_name}")
|
||
|
|
lines.append("自动生成，请勿手动修改")
|
||
|
|
lines.append('"""')
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
seen_models: set[str] = set()
|
||
|
|
for route in routes:
|
||
|
|
rname = route["name"]
|
||
|
|
method = route["method"]
|
||
|
|
if route.get("parsed_params"):
|
||
|
|
seen_models.add(_model_name(rname, method, "Request"))
|
||
|
|
if route.get("parsed_response"):
|
||
|
|
seen_models.add(_model_name(rname, method, "Response"))
|
||
|
|
|
||
|
|
response_models = [m for m in seen_models if m.endswith("Response")]
|
||
|
|
if response_models:
|
||
|
|
plugin_title = plugin_name.title().replace("-", "").replace("_", "")
|
||
|
|
seen_models.add(f"{plugin_title}Response")
|
||
|
|
|
||
|
|
for model_name in sorted(seen_models):
|
||
|
|
lines.append(f"from .pydantic import {model_name}")
|
||
|
|
|
||
|
|
lines.append("")
|
||
|
|
lines.append("__all__ = [")
|
||
|
|
for model_name in sorted(seen_models):
|
||
|
|
lines.append(f' "{model_name}",')
|
||
|
|
lines.append("]")
|
||
|
|
lines.append("")
|
||
|
|
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
|
||
|
|
def scan_config_dir(config_dir: Path) -> list[dict[str, Any]]:
|
||
|
|
"""扫描配置目录，解析所有 .toml 文件中的路由定义"""
|
||
|
|
import tomli
|
||
|
|
|
||
|
|
routes = []
|
||
|
|
toml_files = sorted(config_dir.glob("*.toml"))
|
||
|
|
|
||
|
|
group_config = {}
|
||
|
|
config_file = config_dir / "config.toml"
|
||
|
|
if config_file.exists():
|
||
|
|
with open(config_file, "rb") as f:
|
||
|
|
group_config = tomli.load(f)
|
||
|
|
|
||
|
|
defaults = group_config.get("worker", {}).get("defaults", {})
|
||
|
|
router_prefix = (
|
||
|
|
group_config.get("worker", {})
|
||
|
|
.get("registration", {})
|
||
|
|
.get("router_prefix", "")
|
||
|
|
)
|
||
|
|
|
||
|
|
for toml_file in toml_files:
|
||
|
|
if toml_file.name == "config.toml":
|
||
|
|
continue
|
||
|
|
|
||
|
|
with open(toml_file, "rb") as f:
|
||
|
|
data = tomli.load(f)
|
||
|
|
|
||
|
|
route_name = toml_file.stem
|
||
|
|
|
||
|
|
for method in ("GET", "POST", "PUT", "DELETE", "PATCH", "MQTT"):
|
||
|
|
if method not in data:
|
||
|
|
continue
|
||
|
|
|
||
|
|
entry = data[method]
|
||
|
|
params_raw = entry.get("params", [])
|
||
|
|
response_raw = entry.get("response_model", {})
|
||
|
|
|
||
|
|
route = {
|
||
|
|
"name": route_name,
|
||
|
|
"method": method,
|
||
|
|
"summary": entry.get("summary", ""),
|
||
|
|
"docstring": entry.get("docstring", ""),
|
||
|
|
"internal": entry.get("internal", defaults.get("internal", False)),
|
||
|
|
"params": params_raw,
|
||
|
|
"response_model": response_raw,
|
||
|
|
"parsed_params": parse_params(params_raw),
|
||
|
|
"parsed_response": parse_response_model(response_raw),
|
||
|
|
"path": entry.get("path", f"{router_prefix}/{route_name}"),
|
||
|
|
"timeout": entry.get("timeout", defaults.get("timeout", 3.0)),
|
||
|
|
"auth_required": entry.get("auth_required", defaults.get("auth_required", False)),
|
||
|
|
}
|
||
|
|
routes.append(route)
|
||
|
|
|
||
|
|
return routes
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
if len(sys.argv) < 3:
|
||
|
|
print("用法: uv run gen_models <config_dir> <plugin_dir>")
|
||
|
|
print("示例:")
|
||
|
|
print(" uv run gen_models config/workers/example src/plugins/example")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
config_dir = Path(sys.argv[1])
|
||
|
|
plugin_dir = Path(sys.argv[2])
|
||
|
|
plugin_name = plugin_dir.name
|
||
|
|
|
||
|
|
if not config_dir.exists():
|
||
|
|
print(f"❌ 配置目录不存在: {config_dir}")
|
||
|
|
sys.exit(1)
|
||
|
|
if not plugin_dir.exists():
|
||
|
|
print(f"❌ 插件目录不存在: {plugin_dir}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
print(f"🔍 扫描配置目录: {config_dir}")
|
||
|
|
routes = scan_config_dir(config_dir)
|
||
|
|
print(f" 发现 {len(routes)} 个路由定义")
|
||
|
|
|
||
|
|
if not routes:
|
||
|
|
print("⚠️ 未找到路由定义，跳过模型生成")
|
||
|
|
return
|
||
|
|
|
||
|
|
models_dir = plugin_dir / "models"
|
||
|
|
models_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
pydantic_code = generate_pydantic_code(plugin_name, routes)
|
||
|
|
pydantic_file = models_dir / "pydantic.py"
|
||
|
|
pydantic_file.write_text(pydantic_code, encoding="utf-8")
|
||
|
|
print(f"✅ 生成: {pydantic_file}")
|
||
|
|
|
||
|
|
init_code = generate_init_code(plugin_name, routes)
|
||
|
|
init_file = models_dir / "__init__.py"
|
||
|
|
init_file.write_text(init_code, encoding="utf-8")
|
||
|
|
print(f"✅ 生成: {init_file}")
|
||
|
|
|
||
|
|
print("\n📋 生成的模型:")
|
||
|
|
for route in routes:
|
||
|
|
method = route["method"]
|
||
|
|
name = route["name"]
|
||
|
|
params = route.get("parsed_params", [])
|
||
|
|
response = route.get("parsed_response", [])
|
||
|
|
if params:
|
||
|
|
req_model = _model_name(name, method, "Request")
|
||
|
|
param_strs = [f"{p['name']}: {p['type']}" for p in params]
|
||
|
|
print(f" 📥 {req_model}: {', '.join(param_strs)}")
|
||
|
|
if response:
|
||
|
|
resp_model = _model_name(name, method, "Response")
|
||
|
|
field_strs = [f"{f['name']}: {f['type']}" for f in response]
|
||
|
|
print(f" 📤 {resp_model}: {', '.join(field_strs)}")
|
||
|
|
|
||
|
|
|
||
|
|
def gen_models():
|
||
|
|
"""CLI 入口"""
|
||
|
|
main()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|