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()
|