Скрипт упрощающий работу с `protoc`

grpc_generator.py

"""
Скрипт группирует, разрешает импорты и генерирует клиентский и серверный код из ваших *.proto файлов.
Зависмости: pip install grpcio grpcio-tools

Пример использования: python grpc_generator.py -p protos -o src/grpc
Результат работы этой команды:
- Получит список файлов с расширением *.proto из папки protos
Например: ['user_service.proto', 'auth_service.proto']
- Создаст в директории src директорию grpc
- Создаст директории для каждого *.proto файла в директории src/grpc и создаст в них __init__.py:
--src
  └── grpc
      ├── user_service
      │   └── __init__.py
      └── auth_service
          └── __init__.py

- Сгенерирует файлы *.pb2_grpc.py и *.pb2.py для каждого *.proto файла из директории protos и
сохранит их в директории src/grpc:
--src
  └── grpc
      ├── user_service
      |   ├── __init__.py
      |   └── user_service_pb2.py
      │   └── user_service_pb2_grpc.py
      └── auth_service
          ├── __init__.py
          |── auth_service_pb2.py
          └── auth_service_pb2_grpc.py
"""

from pathlib import Path
from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter

from grpc_tools.protoc import main as protoc
from grpc_tools.protoc import _PROTO_MODULE_SUFFIX, _SERVICE_MODULE_SUFFIX


def view(
    *args, 
    begin_block: bool = False, 
    end_block: bool = False, 
    left: int = None,
    up: int = None, 
    down: int = 1, 
    view_if: bool = True,
    **kwargs
):
    if view_if:
        if up:
            for _ in range(up):
                print('│')

        if begin_block:
            print('\n┌───')
            
        if not begin_block and not end_block:
            left_padding = '─' * left if left else ''
            print('├' + left_padding, *args, **kwargs)

        if end_block:
            print('└───\n')

        if down:
            for _ in range(down):
                print('│')


def args_resolve() -> Namespace:
    parser = ArgumentParser(
        description=__doc__, 
        formatter_class=RawDescriptionHelpFormatter,
        usage='python grpc_generator.py -p [Папка с .proto файлами] -o [Путь до директории, где создать файлы]',
        add_help=True,
    )
    
    main_section = parser.add_argument_group('Основные параметры')
    main_section.add_argument('--no-autoresolve-imports', action='store_true', help='Не выполнять автоматическое исправление импортов в *.py и *.pyi файлах')
    main_section.add_argument('--just-display', action='store_true', help='Только отобразить полученные команды, но не выполнять их')

    paths_section = parser.add_argument_group('Директории и пути')
    paths_section.add_argument('-p', '--proto_dir', type=str, help='Путь до папки с *.proto файлами', required=True)
    paths_section.add_argument('-gpo', '--grpc_python_out', type=str, help='Путь до папки, где будут созданы *pb2_grpc.py файлы. По умолчанию создаст в текущей директории в папке "grpc"')
    paths_section.add_argument('-po', '--python_out', type=str, help='Путь до папки, где будут созданы *pb2.py файлы. По умолчанию создаст в текущей директории в папке "grpc"')
    paths_section.add_argument('-pyio', '--pyi_out', type=str, help='Путь до папки, где будут созданы *.pyi файлы. По умолчанию создаст в текущей директории в папке "grpc"')
    paths_section.add_argument('-o', '--output_dir', type=str, help='Указывает общую папку для аргументов "--python_out", "--grpc_python_out", "--pyi_out"')

    other_section = parser.add_argument_group('Другие параметры')
    other_section.add_argument('-s', '--silent', action='store_true', help='Не показывать подробный вывод')
    other_section.add_argument('-noinit', action='store_true', help='Не создавать "__init__.py" в папках')
    other_section.add_argument('-no-pyi', action='store_true', help='Не создавать *.pyi файлы')
    other_section.add_argument('-no-grpc', action='store_true', help='Не создавать *.pb2_grpc.py файлы')

    return parser.parse_args()


args = args_resolve()
VIEW_LOG = False if args.silent else True

generated_pb2_files = []
generated_pb2_grpc_files = []
generated_pyi_files = []


def search_proto_files() -> list[Path]:
    if not args.proto_dir:
        raise ValueError('Не указан путь до папки с *.proto файлами')
    else:
        proto_dir_path = Path(args.proto_dir)
        view(f'Получен путь {proto_dir_path.resolve()}', view_if=VIEW_LOG)

        proto_files = list(proto_dir_path.glob('*.proto'))
        view(f'Получен список *.proto файлов: {len(proto_files)}', view_if=VIEW_LOG)

        return proto_files


def resolve_out_dirs():
    python_out = './grpc'
    grpc_python_out = './grpc'
    pyi_out = './grpc'
    
    if not args.output_dir:
        if args.python_out:
            python_out = args.python_out
        if args.grpc_python_out:
            grpc_python_out = args.grpc_python_out
        if args.pyi_out:
            pyi_out = args.pyi_out
    else:
        python_out = args.output_dir
        grpc_python_out = args.output_dir
        pyi_out = args.output_dir
    
    view(f'Получен путь для *.py файлов:  {python_out}', view_if=VIEW_LOG)
    view(f'Получен путь для *.pb2_grpc.py файлов:  {grpc_python_out}', view_if=VIEW_LOG)
    view(f'Получен путь для *.pyi файлов:  {pyi_out}', view_if=VIEW_LOG)
    
    return python_out, grpc_python_out, pyi_out


def make_directory(path):
    path_ = Path(path)
    if not args.just_display:
        if not path_.exists():
            view(f'Создание папки {path_.resolve()}', view_if=VIEW_LOG)
            path_.mkdir(parents=True, exist_ok=True)

            return path_


def make_init_file(path):
    path_ = Path(f'{path}/__init__.py')
    if not args.noinit and not args.just_display:
        if not path_.exists():
            view(f'Создание __init__.py в {path_.resolve()}', view_if=VIEW_LOG)
            path_.touch(exist_ok=True)


def get_module_path(filename, no_suffix=False):
    python_out, _, _ = resolve_out_dirs()
    return f'{python_out}/{filename}/{filename}{'' if no_suffix else _PROTO_MODULE_SUFFIX}.py'


def get_service_path(filename, no_suffix=False):
    _, grpc_python_out, _ = resolve_out_dirs()
    return f'{grpc_python_out}/{filename}/{filename}{'' if no_suffix else _SERVICE_MODULE_SUFFIX}.py'


def get_stub_path(filename, no_suffix=False):
    _, _, pyi_out = resolve_out_dirs()
    return f'{pyi_out}/{filename}/{filename}{'' if no_suffix else _PROTO_MODULE_SUFFIX}.pyi'


def resolve_imports():
    # TODO: реализовать функцию, которая разрешает импорты в сгенерированном коде
    
    view(f'Анализирую зависимости... (imports)', view_if=VIEW_LOG)
    
    def read_file(path):
        with open(path, 'r') as file:
            lines = file.readlines()
        return list(map(lambda x: x.strip(), lines)), lines
    
    def write_file(path, lines):
        with open(path, 'w') as file:
            file.writelines(lines)
    
    def search_imports(in_lines):
        for lineno, line in enumerate(in_lines):
            if line.startswith('import'):
                module_name = line.split()[1]
                yield module_name, lineno, line
            
    
    def walk(files: list, pb2_files: list):
        root_dir = Path.cwd()
        
        for file in files:
            view(f'Обработка файла {file.resolve()}', left=2, view_if=VIEW_LOG)
            
            stripped, raw = read_file(file)
            for module_name, lineno, line in search_imports(stripped):
                for pb2 in pb2_files:
                    filename = pb2.stem
                    dependant = Path(pb2).resolve() if module_name == filename else None
                    is_relative = dependant.is_relative_to(root_dir) if dependant else False
                    
                    if is_relative:
                        relative_path = dependant.relative_to(root_dir).with_suffix('')
                        import_path = str(relative_path).replace('/', '.')
                        
                        view(f'Найдена зависимость "{module_name}" (строка {lineno + 1}): {file} --> {pb2}', left=4, view_if=VIEW_LOG)
                        
                        previous_import_line = line
                        new_import_line = line.replace(f'import {module_name}', f'import {import_path}')
                        raw[lineno] = f'{new_import_line}\n'
                        
                        view(f'Импорт изменен: "{previous_import_line}" --> "{new_import_line}"', left=4, view_if=VIEW_LOG)
            
            write_file(file, raw)

    walk(generated_pb2_files, generated_pb2_files)
    walk(generated_pb2_grpc_files, generated_pb2_files)
    walk(generated_pyi_files, generated_pb2_files)


def generate_client(from_proto_file):
    """
    TODO: Доделать функцию генерации клиента из *.proto файла
    """
    import re
    
    view("Я не доделал этот функционал:", up=2, view_if=VIEW_LOG)
    view(f'Генерирую клиент из {from_proto_file}', view_if=VIEW_LOG)
    filename = from_proto_file.stem
    
    def camel_case(s: str) -> str:
        return ''.join(word.capitalize() for word in s.split('_'))
    
    def extract_service_name(path):
        with open(path, 'r') as file:
            lines = file.readlines()
        text = ''.join(lines)
        # Регулярное выражение для поиска строк типа "service UserService {"
        pattern = r'service\s+(\w+)\s*\{'
        matches = re.findall(pattern, text)
        
        return matches  # Возвращает список найденных имен сервисов
    
    services = extract_service_name(from_proto_file)
    view(f'Найдены сервисы: {services}', view_if=VIEW_LOG)
    
    service_module = str(Path(get_service_path(filename)).with_suffix('')).replace('/', '.')
    service_module_stub = camel_case(filename) + 'Stub'
    
    template = f"""
    import grpc
    from {service_module} import {service_module_stub}
    ...
    """
    
    view(f'Сгенерированный код: {template}', view_if=VIEW_LOG)


def main():
    proto_files: list[Path] = search_proto_files()
    
    if not proto_files:
        raise ValueError('*.proto файлы не найдены')

    python_out, grpc_python_out, pyi_out = resolve_out_dirs()

    commands = ['grpc_tools.protoc', f'-I{args.proto_dir}']

    for proto_file in proto_files:
        cmds = list(commands)
        filename = proto_file.stem

        python_out_dir = f'{python_out}/{filename}'
        if make_directory(python_out_dir):
            make_init_file(python_out_dir)
        cmds += [f'--python_out={str(python_out_dir)}']

        if not args.no_grpc:
            grpc_python_out_dir = f'{grpc_python_out}/{filename}'
            if make_directory(grpc_python_out_dir):
                make_init_file(grpc_python_out_dir)
            cmds += [f'--grpc_python_out={str(grpc_python_out_dir)}']

        if not args.no_pyi:
            pyi_out_dir = f'{pyi_out}/{filename}'
            make_directory(pyi_out_dir)
            cmds += [f'--pyi_out={str(pyi_out_dir)}']

        cmds += [str(proto_file.name)]

        result_command = ' '.join(cmds)
        view(f'COMMAND[{filename}]:  {result_command}', left=4, view_if=VIEW_LOG)

        if not args.just_display:
            protoc(cmds)
            
            generated_module = Path(get_module_path(filename))
            generated_service = Path(get_service_path(filename))
            generated_stub = Path(get_stub_path(filename))
            
            if Path.exists(generated_module):
                view(f'Создан файл (pb2)', generated_module, view_if=VIEW_LOG)
                generated_pb2_files.append(generated_module)

            if Path.exists(generated_service):
                view(f'Создан файл (grpc)', generated_service, view_if=VIEW_LOG)
                generated_pb2_grpc_files.append(generated_service)
            
            if Path.exists(generated_stub):
                view(f'Создан файл (stub)', generated_stub, view_if=VIEW_LOG)
                generated_pyi_files.append(generated_stub)

    if not args.no_autoresolve_imports:
        resolve_imports()
    generate_client(proto_file)


if __name__ == '__main__':
    view(begin_block=True, down=0, view_if=VIEW_LOG)
    main()
    view(end_block=True, down=0, view_if=VIEW_LOG)