150 lines
3.8 KiB
Python
150 lines
3.8 KiB
Python
import os
|
|
from typing import Dict, List
|
|
import re
|
|
from jinja2.environment import Template
|
|
|
|
sysPath = os.getcwd()
|
|
protoAddress = f"{sysPath}/protobuf"
|
|
tempPath = f"{sysPath}/storage/temp"
|
|
daoPath = f"{sysPath}/storage/dao"
|
|
servicePath = f"{sysPath}/storage/service"
|
|
|
|
|
|
class ProtoOption(object):
|
|
def __init__(self, method: str) -> None:
|
|
super().__init__()
|
|
self.method = method
|
|
self.option = []
|
|
self.ret = []
|
|
|
|
def add_option(self, opt: List[str]):
|
|
self.option.extend(opt)
|
|
|
|
def add_ret(self, opt: List[str]):
|
|
self.ret.extend(opt)
|
|
|
|
def __str__(self) -> str:
|
|
req = ", ".join([f"{i[0]} {i[1]}" for i in self.option])
|
|
resp = ", ".join([f"{i[0]} {i[1]}" for i in self.ret])
|
|
return f"{self.method}({req})({resp})"
|
|
|
|
|
|
def dist_to_ProOpt(req, resp) -> List[ProtoOption]:
|
|
|
|
def parse_type(l: str) -> List[str]:
|
|
l = l.strip()
|
|
if l == "":
|
|
return []
|
|
|
|
opt = l.split(";")
|
|
result = []
|
|
for l_opt in opt:
|
|
l_opt = l_opt.strip()
|
|
l_list = l_opt.split()
|
|
if len(l_list) == 0:
|
|
continue
|
|
|
|
val = l_list[0]
|
|
if val == "BaseKey":
|
|
val = "*proto.BaseKey"
|
|
result.append([l_list[1].capitalize(), val])
|
|
|
|
elif val == "repeated":
|
|
val = f"[]{l_list[1]}"
|
|
result.append([l_list[2].capitalize(), val])
|
|
|
|
else:
|
|
result.append([l_list[1].capitalize(), val])
|
|
return result
|
|
|
|
lists = []
|
|
for key, value in req.items():
|
|
p = ProtoOption(method=key)
|
|
p.add_option(parse_type(value))
|
|
p.add_ret(parse_type(resp.get(key, "")))
|
|
lists.append(p)
|
|
return lists
|
|
|
|
|
|
def parse_protobuf_to_Opt(name: str) -> List[ProtoOption]:
|
|
|
|
gather_req = re.compile(r'^(.*?)Request\s+{(.*?)}$')
|
|
gather_resp = re.compile(r'^(.*?)Response\s+{(.*?)}$')
|
|
|
|
with open(f"{protoAddress}/{name}", "r", encoding="utf-8") as fp:
|
|
text = fp.read()
|
|
text = re.sub(r'//.+', "", text) # 去掉注释
|
|
text = text.replace("\n", "")
|
|
result = text.split("message")[1:]
|
|
|
|
reqDist = {}
|
|
respDist = {}
|
|
|
|
for r in result:
|
|
req = gather_req.findall(r)
|
|
resp = gather_resp.findall(r)
|
|
if len(req) != 0:
|
|
req = req[0]
|
|
reqDist[req[0].strip()] = req[1].strip()
|
|
if len(resp) != 0:
|
|
resp = resp[0]
|
|
respDist[resp[0].strip()] = resp[1].strip()
|
|
|
|
return dist_to_ProOpt(reqDist, respDist)
|
|
|
|
|
|
def load_protobuf() -> List[ProtoOption]:
|
|
for _, _, file_name_list in os.walk(protoAddress):
|
|
break
|
|
|
|
li = []
|
|
for name in file_name_list:
|
|
if name.endswith("x.proto"):
|
|
li.extend(parse_protobuf_to_Opt(name))
|
|
return li
|
|
|
|
def go_fmt(path: str):
|
|
os.system(f"go fmt {path}")
|
|
|
|
|
|
def load_template(name: str) -> str:
|
|
with open(f"{tempPath}/{name}", "r", encoding="utf-8") as fp:
|
|
return fp.read()
|
|
|
|
def gen_dao_interface(proto):
|
|
|
|
tem_text = load_template("dao.template")
|
|
template = Template(tem_text)
|
|
|
|
text = template.render(keys=proto)
|
|
|
|
temp_path = f"{daoPath}/interface.gen.go"
|
|
with open(temp_path, 'w', encoding='utf-8') as f:
|
|
f.write(text)
|
|
|
|
|
|
def gen_single_service(proto):
|
|
tem_text = load_template("service.template")
|
|
template = Template(tem_text)
|
|
|
|
text = template.render(keys=proto)
|
|
|
|
temp_path = f"{servicePath}/single_service.gen.go"
|
|
with open(temp_path, 'w', encoding='utf-8') as f:
|
|
f.write(text)
|
|
|
|
def format_code_go():
|
|
go_fmt(f"{daoPath}/interface.gen.go")
|
|
go_fmt(f"{servicePath}/single_service.gen.go")
|
|
|
|
def main():
|
|
# 加载 protobuf
|
|
protobuf = load_protobuf()
|
|
# 生成 dao 接口
|
|
gen_dao_interface(protobuf)
|
|
gen_single_service(protobuf)
|
|
format_code_go()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|