wheat-cache/shell/gen_protobuf.py

import os

import yaml

sysPath = os.getcwd()
tempPath = f"{sysPath}/storage/temp"
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, Loader=yaml.FullLoader)

    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):
    for key, value in cfg_camel.items():
        proto_path = f"{protobufPath}/{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}\n')
                file.write('\nmessage ' + v + 'Response ' + '{\n}\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}\n')
                    file.write("\nmessage %sResponse {\n}\n" % v)
                    file.close()

        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('\n\nservice CommServer {\n')
    for key, value in cfg_camel.items():
        for v in value:
            file.write('  rpc ' + v + ' (' + v + 'Request) returns (%s);\n' % "{}Response".format(v))
    file.write('}')
    file.close()
    print("storage.proto", "-> success")


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