python 生成 proto 文件

import os
import sys
import subprocess
import shutil


"""
    功能：
        1. 拷貝python的proto到go的對應目錄之下
        2. 生成python的源碼 - import .
        3. 生成go的源碼
"""

class cd:
    def __init__(self, newPath):
        self.newPath = os.path.expanduser(newPath)

    def __enter__(self):
        self.savedPath = os.getcwd()
        os.chdir(self.newPath)

    def __exit__(self, etype, value, traceback):
        os.chdir(self.savedPath)


def replace_file(file_name):
    new_file_name = f"{file_name}-bak"
    modify_times = 0   # 統計修改次數
    f1 = open(file_name,'r',encoding='utf-8')    # 以“r”(只讀)模式打開舊文件
    f2 = open(new_file_name,'w',encoding='utf-8')  # 以“w”(寫)模式打開或創建新文件（寫模式下，文件存在則重寫文件，文件不存在則創建文件）
    for lines in f1: # 循環讀
        if lines.startswith("import") and not lines.startswith("import grpc"):
            lines = f"from . {lines}"
            modify_times += 1  # 每修改一次就自增一次
        f2.write(lines) # 將修改后的內容后的內容寫入新文件
    print('文件修改的次數：',modify_times) # 9
    f1.close()  # 關閉文件f1
    f2.close()  # 關閉文件f2(同時打開多個文件時，先打開的先關閉，后打開的后關閉)
    os.replace(new_file_name,file_name) # 修改(替換)文件名


def proto_file_list(path):
    flist = []
    lsdir = os.listdir(path)
    dirs = [i for i in lsdir if os.path.isdir(os.path.join(path, i))]
    if dirs:
        for i in dirs:
            proto_file_list(os.path.join(path, i))
    files = [i for i in lsdir if os.path.isfile(os.path.join(path, i))]
    for file in files:
        if file.endswith(".proto"):
            flist.append(file)
    return flist


def copy_from_py_to_go(src_dir, dst_dir):
    proto_files = proto_file_list(src_dir)
    for proto_file in proto_files:
        try:
            shutil.copy(f"{src_dir}/{proto_file}", dst_dir)
        except IOError as e:
            print("Unable to copy file. %s" % e)
        except:
            print("Unexpected error:", sys.exc_info())


def generated_pyfile_list(path):
    flist = []
    lsdir = os.listdir(path)
    dirs = [i for i in lsdir if os.path.isdir(os.path.join(path, i))]
    if dirs:
        for i in dirs:
            proto_file_list(os.path.join(path, i))
    files = [i for i in lsdir if os.path.isfile(os.path.join(path, i))]
    for file in files:
        if file.endswith(".py"):
            flist.append(file)
    return flist


class ProtoGenerator:
    def __init__(self, python_dir, go_dir):
        self.python_dir = python_dir
        self.go_dir = go_dir

    def generate(self):
        with cd(self.python_dir):
            files = proto_file_list(self.python_dir)
            subprocess.call("workon mxshop_srv", shell=True)
            for file in files:
                command = f"python -m grpc_tools.protoc --python_out=. --grpc_python_out=. -I. {file}"
                subprocess.call(command)

            #查詢生成的py文件並添加上 from .
            py_files = generated_pyfile_list(self.python_dir)
            for file_name in py_files:
                replace_file(file_name)
            print(py_files)

        with cd(self.go_dir):
            files = proto_file_list(self.go_dir)
            for file in files:
                command = f"protoc -I . {file} --go_out=plugins=grpc:."
                subprocess.call(command)


class PyProtoGenerator:
    def __init__(self, python_dir):
        self.python_dir = python_dir

    def generate(self):
        with cd(self.python_dir):
            files = proto_file_list(self.python_dir)
            subprocess.call("workon mxshop_srv", shell=True)
            for file in files:
                command = f"python -m grpc_tools.protoc --python_out=. --grpc_python_out=. -I. {file}"
                subprocess.call(command)

            #查詢生成的py文件並添加上 from .
            py_files = generated_pyfile_list(self.python_dir)
            for file_name in py_files:
                replace_file(file_name)
            print(py_files)


if __name__ == "__main__":
    #goods的proto生成
    # python_dir = "C:\\Users\\Administrator\\PycharmProjects\\mxshop_srvs\\userop_srv\\proto"
    # go_dir = "D:\\python微服務\\projects\\mxshop-api\\userop-web\\proto"
    # #
    # # #將py目錄下的文件夾拷貝到go目錄下
    # copy_from_py_to_go(python_dir, go_dir)
    # #
    # # #生成對應的py源碼和go源碼
    # gen = ProtoGenerator(python_dir, go_dir)
    # gen.generate()

    py_gen = PyProtoGenerator("C:\\Users\\Administrator\\PycharmProjects\\mxshop_srvs\\inventory_srv\\proto")
    py_gen.generate()