101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
自动生成 gRPC Python 代码脚本（相对导入版）
|
|||
|
|
递归扫描 protos/ 目录并生成到 generated/
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import subprocess
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_grpc_code(proto_dir: str, output_dir: str):
|
|||
|
|
proto_path = Path(proto_dir)
|
|||
|
|
if not proto_path.exists():
|
|||
|
|
raise FileNotFoundError(f"❌ proto 目录不存在: {proto_dir}")
|
|||
|
|
|
|||
|
|
proto_files = [str(p) for p in proto_path.rglob("*.proto")]
|
|||
|
|
if not proto_files:
|
|||
|
|
print("⚠️ 未找到任何 .proto 文件")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
os.makedirs(output_dir, exist_ok=True)
|
|||
|
|
print(f"🔧 生成目录: {output_dir}")
|
|||
|
|
|
|||
|
|
for proto_file in proto_files:
|
|||
|
|
print(f" ⏳ 编译 {proto_file}")
|
|||
|
|
|
|||
|
|
# 只需要指定根目录
|
|||
|
|
include_paths = [f"-I{proto_dir}"]
|
|||
|
|
|
|||
|
|
cmd = [
|
|||
|
|
"python",
|
|||
|
|
"-m",
|
|||
|
|
"grpc_tools.protoc",
|
|||
|
|
*include_paths,
|
|||
|
|
f"--python_out={output_dir}",
|
|||
|
|
f"--grpc_python_out={output_dir}",
|
|||
|
|
proto_file,
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
result = subprocess.run(cmd, check=True, capture_output=True, text=True)
|
|||
|
|
if result.stdout.strip():
|
|||
|
|
print(result.stdout)
|
|||
|
|
if result.stderr.strip():
|
|||
|
|
print("⚠️ 编译警告:", result.stderr)
|
|||
|
|
except subprocess.CalledProcessError as e:
|
|||
|
|
print("❌ 编译失败:")
|
|||
|
|
print(e.stdout)
|
|||
|
|
print(e.stderr)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 确保每个目录都有 __init__.py
|
|||
|
|
ensure_init_files(output_dir)
|
|||
|
|
|
|||
|
|
# 修改 _pb2_grpc.py 为相对导入
|
|||
|
|
convert_grpc_imports_to_relative(output_dir)
|
|||
|
|
|
|||
|
|
print(f"✅ 所有 proto 文件已生成到: {output_dir}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def ensure_init_files(output_dir: str):
|
|||
|
|
"""递归生成 __init__.py 确保 Python 包可导入"""
|
|||
|
|
for root, dirs, files in os.walk(output_dir):
|
|||
|
|
init_file = os.path.join(root, "__init__.py")
|
|||
|
|
if not os.path.exists(init_file):
|
|||
|
|
open(init_file, "w", encoding="utf-8").close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def convert_grpc_imports_to_relative(output_dir: str):
|
|||
|
|
"""把 _pb2_grpc.py 中的绝对导入改为相对导入"""
|
|||
|
|
for root, dirs, files in os.walk(output_dir):
|
|||
|
|
for file in files:
|
|||
|
|
if file.endswith("_pb2_grpc.py"):
|
|||
|
|
path = os.path.join(root, file)
|
|||
|
|
with open(path, "r", encoding="utf-8") as f:
|
|||
|
|
lines = f.readlines()
|
|||
|
|
new_lines = []
|
|||
|
|
for line in lines:
|
|||
|
|
# 将类似 `import xxx_pb2 as xxx__pb2` 改为 `from . import xxx_pb2 as xxx__pb2`
|
|||
|
|
if line.startswith("import ") and "_pb2 as " in line:
|
|||
|
|
parts = line.split(" as ")
|
|||
|
|
module_name = parts[0].replace("import", "").strip()
|
|||
|
|
alias = parts[1].strip()
|
|||
|
|
new_lines.append(f"from . import {module_name} as {alias}\n")
|
|||
|
|
else:
|
|||
|
|
new_lines.append(line)
|
|||
|
|
with open(path, "w", encoding="utf-8") as f:
|
|||
|
|
f.writelines(new_lines)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
ROOT_DIR = Path(__file__).resolve().parents[1]
|
|||
|
|
PROTO_DIR = ROOT_DIR / "protos"
|
|||
|
|
OUTPUT_DIR = ROOT_DIR / "generated"
|
|||
|
|
|
|||
|
|
print("🚀 开始生成 gRPC Python 代码（相对导入版）...")
|
|||
|
|
generate_grpc_code(str(PROTO_DIR), str(OUTPUT_DIR))
|
|||
|
|
print("\n🎉 gRPC Python 代码生成完成！")
|