321 lines
10 KiB
Python
321 lines
10 KiB
Python
"""
|
||
插件脚手架生成器 — 生成风格统一的 Worker 插件模板
|
||
|
||
用法:
|
||
uv run gen-plugin <name> [operation ...]
|
||
|
||
示例:
|
||
# 生成插件骨架（含一个示例操作）:
|
||
uv run gen-plugin myplugin
|
||
|
||
# 指定初始操作:
|
||
uv run gen-plugin myplugin analyze count
|
||
|
||
# 为已存在的插件增量添加操作:
|
||
uv run gen-plugin myplugin add
|
||
|
||
工作流:
|
||
1. uv run gen-plugin myplugin # 生成 config + src 骨架（统一风格）
|
||
2. uv run gen_models config/workers/myplugin src/plugins/myplugin
|
||
# 生成 Pydantic 模型
|
||
3. 编辑 src/plugins/myplugin/router/<operation>.py 实现业务逻辑
|
||
4. 编辑 src/plugins/myplugin/__init__.py 的 on_init / on_cleanup
|
||
5. uv run serve # 启动验证
|
||
|
||
生成内容:
|
||
config/workers/<name>/config.toml Worker 注册 + 运行时配置
|
||
config/workers/<name>/<op>.toml 每个操作的 RPC 端点定义
|
||
src/plugins/<name>/__init__.py 基于 Plugin 基类的统一入口
|
||
src/plugins/<name>/global_var.py 白名单版全局变量容器
|
||
src/plugins/<name>/router/__init__.py 端点模块聚合导出
|
||
src/plugins/<name>/router/<op>.py 每个操作的 RPC 实现骨架
|
||
src/plugins/<name>/models/ 由 gen_models 生成（见步骤 2）
|
||
"""
|
||
import re
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from ._root import find_project_root
|
||
|
||
# 项目根目录（CHONGMING_ROOT 环境变量 > CWD 向上查找，见 _root.py）
|
||
ROOT = find_project_root()
|
||
CONFIG_DIR = ROOT / "config" / "workers"
|
||
PLUGIN_DIR = ROOT / "src" / "plugins"
|
||
|
||
NAME_RE = re.compile(r"^[a-z][a-z0-9_]*$")
|
||
|
||
|
||
# ============================================================
|
||
# 模板
|
||
# ============================================================
|
||
|
||
def snake_to_camel(name: str) -> str:
|
||
"""snake_case -> CamelCase（如 analyze_media -> AnalyzeMedia）"""
|
||
return "".join(part.title() for part in name.split("_"))
|
||
|
||
|
||
def _gen_config_toml(name: str) -> str:
|
||
return f'''# {name} Worker 插件配置
|
||
|
||
[worker]
|
||
[worker.defaults]
|
||
timeout = 30
|
||
cacheable = false
|
||
|
||
[worker.registration]
|
||
router_prefix = "/{name}"
|
||
|
||
# Worker 注册（app lifespan 据此发现并加载插件）
|
||
[workers]
|
||
type = "{name}"
|
||
module = "plugins.{name}"
|
||
factory = "create_workers"
|
||
cleanup = "cleanup"
|
||
count = 1
|
||
|
||
# 运行时配置（展平后传给 create_workers 的 config 参数）
|
||
[workers.config]
|
||
data_dir = "./data"
|
||
'''
|
||
|
||
|
||
def _gen_operation_toml(name: str, operation: str) -> str:
|
||
return f'''# {name}.{operation} — {operation} 操作
|
||
# 修改后需重新运行: uv run gen_models config/workers/{name} src/plugins/{name}
|
||
[GET]
|
||
cacheable = true
|
||
summary = "{operation} 操作描述"
|
||
params = ["value: str"]
|
||
|
||
[GET.response_model]
|
||
result = {{ type = "str", required = true }}
|
||
'''
|
||
|
||
|
||
def _gen_plugin_init(name: str) -> str:
|
||
cls = snake_to_camel(name)
|
||
return f'''"""{name.capitalize()} Worker 插件
|
||
|
||
统一风格模板 — 基于 chongming_worker.plugin.Plugin 基类。
|
||
|
||
工作流:
|
||
1. 编辑 config/workers/{name}/*.toml 定义 RPC 端点
|
||
2. uv run gen_models config/workers/{name} src/plugins/{name}
|
||
3. 在 router/<操作>.py 中实现业务逻辑
|
||
4. 在 on_init / on_cleanup 中初始化 / 清理资源
|
||
"""
|
||
import logging
|
||
|
||
from chongming_worker.plugin import Plugin
|
||
|
||
from . import router # noqa: F401 触发 @rpc_method 注册
|
||
# from .routes import router # 可选: 自定义 FastAPI 路由（网关自动发现）
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class {cls}Plugin(Plugin):
|
||
"""{name} 插件 — 统一生命周期。"""
|
||
|
||
type = "{name}"
|
||
db_name = "{name}.db" # 自动在 data_dir 下建 SQLite 库（不需要则置 None）
|
||
auto_bind = True # 首次初始化时自动 auto_bind_models(self.type)
|
||
|
||
async def on_init(self):
|
||
"""首次 create_workers 时调用一次：初始化资源、启动后台任务。
|
||
|
||
数据库已由基类 _bootstrap 自动初始化（self.db）。
|
||
"""
|
||
self.set_var("db", self.db)
|
||
|
||
# 需要后台任务时:
|
||
# self.spawn_task(self._background_loop())
|
||
|
||
logger.info("[{name}] 初始化完成")
|
||
|
||
async def on_cleanup(self):
|
||
"""cleanup 时调用：关闭自定义客户端等（后台任务与数据库由基类处理）。"""
|
||
logger.info("[{name}] 清理完成")
|
||
|
||
|
||
# 暴露给 app lifespan 的模块级入口（bound method，签名与旧写法完全一致）
|
||
_plugin = {cls}Plugin()
|
||
create_workers = _plugin.create_workers
|
||
cleanup = _plugin.cleanup
|
||
'''
|
||
|
||
|
||
def _gen_global_var(name: str) -> str:
|
||
return f'''"""{name.capitalize()} 插件全局运行时状态 — 统一容器
|
||
|
||
约定:
|
||
- 模块只做"声明默认值"，不产生任何导入期副作用
|
||
- 所有值统一在 Plugin.on_init 中通过 set_global_var 显式写入
|
||
- set_global_var 仅允许写入 _ALLOWED 白名单，拼写错误会立即抛错
|
||
- require() 用于取"必须已初始化"的值，未初始化时抛清晰错误
|
||
"""
|
||
from typing import Any
|
||
|
||
from chongming_database import Database
|
||
|
||
# 允许写入的全局变量白名单（防止拼写错误静默创建无用的新变量）
|
||
_ALLOWED = frozenset({{"db"}})
|
||
|
||
# 显式声明所有运行时状态（默认 None，运行时由 on_init 写入）
|
||
db: Database | None = None
|
||
|
||
|
||
def set_global_var(name: str, value: Any) -> None:
|
||
"""白名单校验版 setter — 设置模块级运行时状态。"""
|
||
if name not in _ALLOWED:
|
||
raise KeyError(
|
||
f"global_var: 未知变量 '{{name}}'（允许: {{sorted(_ALLOWED)}}）"
|
||
)
|
||
globals()[name] = value
|
||
|
||
|
||
def require(name: str) -> Any:
|
||
"""获取必须已初始化的全局变量，未初始化时抛清晰错误。"""
|
||
if name not in _ALLOWED:
|
||
raise KeyError(f"global_var: 未知变量 '{{name}}'（允许: {{sorted(_ALLOWED)}}）")
|
||
value = globals().get(name)
|
||
if value is None:
|
||
raise RuntimeError(
|
||
f"global_var: '{{name}}' 未初始化 —— 请确认插件 create_workers 已执行且配置正确"
|
||
)
|
||
return value
|
||
|
||
|
||
__all__ = sorted(_ALLOWED) + ["set_global_var", "require"]
|
||
'''
|
||
|
||
|
||
def _gen_router_init(name: str, operations: list[str]) -> str:
|
||
lines = [f'"""{name.capitalize()} RPC 端点"""', ""]
|
||
for op in operations:
|
||
lines.append(f"from . import {op}")
|
||
lines.append("")
|
||
lines.append("__all__ = " + repr(operations))
|
||
lines.append("")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _gen_router_operation(name: str, operation: str) -> str:
|
||
op_camel = snake_to_camel(operation)
|
||
return f'''"""RPC 端点: {name}.{operation}
|
||
|
||
模型由 gen_models 生成（编辑 config/workers/{name}/{operation}.toml 后重新生成）。
|
||
"""
|
||
from chongming_worker.decorators import rpc_method
|
||
|
||
from ..models.pydantic import {op_camel}GetRequest, {op_camel}GetResponse
|
||
|
||
|
||
@rpc_method(
|
||
subject="{name}.{operation}",
|
||
method="GET",
|
||
summary="{operation} 操作",
|
||
response_model={op_camel}GetResponse,
|
||
)
|
||
async def {operation}(request: {op_camel}GetRequest) -> {op_camel}GetResponse:
|
||
"""业务逻辑实现。"""
|
||
return {op_camel}GetResponse(result=f"hello: {{request.value}}")
|
||
'''
|
||
|
||
|
||
# ============================================================
|
||
# 主逻辑
|
||
# ============================================================
|
||
|
||
def _write(path: Path, content: str) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
if path.exists():
|
||
print(f"⏭️ 跳过（已存在）: {path.relative_to(ROOT)}")
|
||
return
|
||
path.write_text(content, encoding="utf-8")
|
||
print(f"✅ 生成: {path.relative_to(ROOT)}")
|
||
|
||
|
||
def _existing_operations(plugin_dir: Path) -> list[str]:
|
||
"""扫描 router/ 下已有的操作文件（排除 __init__.py）。"""
|
||
router_dir = plugin_dir / "router"
|
||
if not router_dir.exists():
|
||
return []
|
||
ops = []
|
||
for f in sorted(router_dir.glob("*.py")):
|
||
if f.name == "__init__.py":
|
||
continue
|
||
ops.append(f.stem)
|
||
return ops
|
||
|
||
|
||
def generate_plugin(name: str, operations: list[str]) -> int:
|
||
"""生成插件骨架。返回 0 表示成功，1 表示失败。"""
|
||
if not NAME_RE.match(name):
|
||
print(f"❌ 插件名不合法: '{name}'（需小写字母/数字/下划线，且以字母开头）")
|
||
return 1
|
||
|
||
config_dir = CONFIG_DIR / name
|
||
plugin_dir = PLUGIN_DIR / name
|
||
is_new = not plugin_dir.exists()
|
||
|
||
# ---- 首次创建: 完整骨架 ----
|
||
if is_new:
|
||
_write(config_dir / "config.toml", _gen_config_toml(name))
|
||
for op in operations:
|
||
_write(config_dir / f"{op}.toml", _gen_operation_toml(name, op))
|
||
_write(plugin_dir / "__init__.py", _gen_plugin_init(name))
|
||
_write(plugin_dir / "global_var.py", _gen_global_var(name))
|
||
(plugin_dir / "models" / "__init__.py").parent.mkdir(parents=True, exist_ok=True)
|
||
(plugin_dir / "models" / "__init__.py").write_text(
|
||
'"""模型导出 — 由 gen_models 生成，请勿手动修改"""\n',
|
||
encoding="utf-8",
|
||
)
|
||
print(f"✅ 生成: {plugin_dir.relative_to(ROOT)}/models/__init__.py")
|
||
|
||
# ---- 增量: 仅补 operations（已存在文件不覆盖）----
|
||
existing = _existing_operations(plugin_dir)
|
||
added = [op for op in operations if op not in existing]
|
||
if added:
|
||
for op in added:
|
||
_write(config_dir / f"{op}.toml", _gen_operation_toml(name, op))
|
||
_write(plugin_dir / "router" / f"{op}.py", _gen_router_operation(name, op))
|
||
|
||
# ---- 重写 router/__init__.py（聚合所有操作）----
|
||
all_ops = _existing_operations(plugin_dir)
|
||
if all_ops:
|
||
init_file = plugin_dir / "router" / "__init__.py"
|
||
init_file.parent.mkdir(parents=True, exist_ok=True)
|
||
init_file.write_text(_gen_router_init(name, all_ops), encoding="utf-8")
|
||
print(f"✅ 更新: {init_file.relative_to(ROOT)}")
|
||
|
||
print(f"\n📋 插件 '{name}' 已就绪，下一步:")
|
||
print(f" 1. uv run gen_models config/workers/{name} src/plugins/{name}")
|
||
print(f" 2. 编辑 src/plugins/{name}/router/ 实现业务逻辑")
|
||
print(" 3. uv run serve")
|
||
return 0
|
||
|
||
|
||
def main() -> None:
|
||
args = sys.argv[1:]
|
||
if not args:
|
||
print("用法: uv run gen_plugin <name> [operation ...]")
|
||
print("示例:")
|
||
print(" uv run gen_plugin myplugin")
|
||
print(" uv run gen_plugin myplugin analyze count")
|
||
print(" uv run gen_plugin myplugin add # 增量添加操作")
|
||
sys.exit(1)
|
||
|
||
name = args[0]
|
||
operations = args[1:] or ["ping"]
|
||
sys.exit(generate_plugin(name, operations))
|
||
|
||
|
||
def gen_plugin():
|
||
"""CLI 入口（供 pyproject.toml [project.scripts] 调用）"""
|
||
main()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|