cmvr-es-cli/generate_proto.py
2025-10-21 09:33:47 +08:00

74 lines
2.2 KiB
Python

import os
import subprocess
from pathlib import Path
def generate_proto():
# 获取当前目录
current_dir = Path(__file__).parent
proto_dir = current_dir / "protos"
generated_dir = current_dir / "generated"
# 创建生成的目录
generated_dir.mkdir(parents=True, exist_ok=True)
# 编译所有 .proto 文件
proto_files = []
for root, _, files in os.walk(proto_dir):
for file in files:
if file.endswith(".proto"):
proto_files.append(os.path.join(root, file))
for proto_path in proto_files:
# 计算相对路径
rel_path = os.path.relpath(proto_path, proto_dir)
output_dir = generated_dir
cmd = [
"python", "-m", "grpc_tools.protoc",
f"-I{proto_dir}",
f"--python_out={output_dir}",
f"--grpc_python_out={output_dir}",
proto_path
]
print(f"Generating code for {rel_path}...")
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
print(f"Error generating {rel_path}:")
print(result.stderr)
return False
print("Proto files generated successfully!")
# 修复生成的代码中的导入路径
fix_imports(generated_dir)
return True
def fix_imports(generated_dir):
"""修复生成的代码中的导入路径"""
for root, _, files in os.walk(generated_dir):
for file in files:
if file.endswith(".py"):
file_path = os.path.join(root, file)
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
# 替换导入路径
content = content.replace(
"from cmvr.api import",
"from generated.cmvr.api import"
)
content = content.replace(
"import cmvr.api.",
"import generated.cmvr.api."
)
with open(file_path, 'w', encoding='utf-8') as f:
f.write(content)
print(f"Fixed imports in {file_path}")
if __name__ == "__main__":
generate_proto()