wheat-cache/shell/gen_protobuf.py

import os
from typing import Dict, List
from jinja2 import Template
import yaml

sysPath = os.getcwd()
tempPath = f"{sysPath}/pkg/structure/generate"
structurePath = f"{sysPath}/pkg/structure"
protobufPath = f"{sysPath}/protobuf"

def to_camel(val: str) -> str:
    return "".join([k.capitalize() for k in val.split('_')])

def load_conf():
    conf_path = f"{tempPath}/tem.yaml"
    with open(conf_path, 'r', encoding='utf-8') as f:
        cfg = f.read()

    cfg = yaml.load(cfg)

    cfg_camel = {}

    for key, val in cfg.items():
        key = to_camel(key).lower()
        cfg_camel[key] = [to_camel(v) for v in val]

    return cfg, cfg_camel


'''生成对应的数据结构proto文件'''
def mk_structure(cfg_camel):
    path = "protobuf"
    storagePath = f"{protobufPath}/storage.proto"
    for key, value in cfg_camel.items():
        proto_path = path + '/' + key + '.proto'
        if not os.path.exists(proto_path):   # 如果这个文件不存在
            '''生成对应的数据结构proto文件'''
            file = open(proto_path,'w')
            file.write('syntax = "proto3";\nimport "base.proto";\noption go_package = "pkg/proto";\n')
            for v in value:
                file.write('\nmessage ' + v +'Request '+'{\n  BaseKey key = 1;\n}')
            file.close()

        else:  # 如果这个文件存在
            '''更新数据结构proto文件'''
            with open(proto_path) as f:
              line = f.readlines()
            f.close()
            for v in value:
                function = v + 'Request'
                flag = 0
                for l in line:
                    if function in l:
                        flag=1
                        break
                if flag == 0:
                    file = open(proto_path,'a')
                    file.write('\nmessage ' + v +'Request '+'{\n  BaseKey key = 1;\n}')
                    file.close()
                flag = 0

        print(f"{key}.proto", "-> success")


def mk_storage(cfg_camel):
    storagePath = f"{protobufPath}/storage.proto"
    file = open(storagePath,'w')
    file.write('// Code generated by gen-struct. DO NOT EDIT.\n// make gen-protobuf generated\n\nsyntax = "proto3";\n\noption go_package = "pkg/proto";\n\n')
    for key,value in cfg_camel.items():
        file.write('import "'+ key +'.proto";\n')
    file.write('\nmessage CommendResponse {\n  repeated string result = 1;\n}\n\nservice CommServer {\n')
    for key,value in cfg_camel.items():
        for v in value:
            file.write('  rpc '+ v +' ('+ v +'Request) returns (CommendResponse);\n')
    file.write('}')
    file.close()
    print("storage.proto", "-> success")


if __name__ == "__main__":
    conf, cfg_camel = load_conf()
    mk_structure(cfg_camel)  # 生成对应的数据结构proto文件
    mk_storage(cfg_camel)