setup.py

import os
import re
import fileinput
import setuptools
import subprocess
from pathlib import Path


package_dir = "helloworld_proto"


def clean_build_directories():
    os.system("rm -rf *.egg-info")
    os.system("rm -rf dist")
    os.system("rm -rf build")
    os.system(f"rm -rf {package_dir}")


def generate_python_classes_from_proto():
    protos_dir = "api"

    os.system(f"mkdir {package_dir}")
    
    python_out = f"./{package_dir}"
    python_grpc_out = f"./{package_dir}"
    inputs = f"./{protos_dir}"
    files = list(map(str, list(Path(f"./{protos_dir}").rglob("*.proto"))))
    files = remove_symlinks_to_protos(files)  # protoc is failing on symlinks to already included protos

    command = [
        "python3", "-m", "grpc_tools.protoc",
        f"-I {inputs}",
        f"--python_out={python_out}",
        f"--grpc_python_out={python_grpc_out}",
        f"--mypy_out={python_out}",
        f"--mypy_grpc_out={python_grpc_out}",
        "--experimental_allow_proto3_optional",
        *files
    ]
    command = " ".join(command)
    
    work_dir = os.getcwd()
    exit_code = subprocess.Popen(command, cwd=work_dir, shell=True).wait()
    if exit_code != 0:
        print("Protobuf generation failed, please check logs above")
        exit(exit_code)


def remove_symlinks_to_protos(proto_files):
    clean_files = []

    for proto_file in proto_files:
        real_path = os.path.realpath(proto_file)
        relative_real = os.path.relpath(real_path)

        if os.path.islink(proto_file) and relative_real in proto_files:
            print(f"Excluding file {proto_file} because it is a symlink for {relative_real}")
        else:
            clean_files.append(proto_file)

    return clean_files


def insert_init_files():
    for root, dirs, files in os.walk(f"./{package_dir}"):
        init_file_path = os.path.join(os.path.abspath(root), "__init__.py")
        with open(init_file_path, "w"):
            pass


def fix_imports():
    def rewrite_py_file(file_input):
        for line in file_input:
            match = re.search("(?<!google\.protobuf) import .*_pb2 as", line)

            if match:
                line = line.replace("from ", f"from {package_dir}.")

            print(line, end="")

    def rewrite_pyi_file(file_input):
        for line in file_input:
            matches = {
                match for match in re.findall(r"[\w\.]*_pb2", line)
                if not match.startswith("google.protobuf")
            }
            for match in matches:
                line = line.replace(match, f"{package_dir}.{match}")

            print(line, end="")

    for root, dirs, files in os.walk(f"./{package_dir}"):
        for file_name in files:
            file_path = os.path.join(os.path.abspath(root), file_name)

            with fileinput.input(file_path, inplace=True) as f:
                if file_name.endswith(".py"):
                    rewrite_py_file(f)
                elif file_name.endswith(".pyi"):
                    rewrite_pyi_file(f)


def setup():
    required = ["googleapis-common-protos>=1.52.0"]

    setuptools.setup(
        name="helloworld-proto",
        use_scm_version={
            "version_scheme": "post-release",
            "local_scheme": "no-local-version",
            "relative_to": __file__,
        },
        setup_requires=["wheel", "setuptools_scm"],
        author="Reface",
        description="Proto example package",
        url="https://github.com/RefaceAI/helloworld-proto",
        packages=setuptools.find_packages(),
        install_requires=required,
        classifiers=[
            "Programming Language :: Python",
            "Programming Language :: Python :: 3",
            "Programming Language :: Python :: 3.8",
            "Programming Language :: Python :: 3.9",
            "Programming Language :: Python :: 3.10",
            "Programming Language :: Python :: 3.11",
            "Programming Language :: Python :: 3.12",
            "Programming Language :: Python :: 3.13",
        ],
        package_data={
            "": ["*.pyi"],
        },
    )


def run():
    clean_build_directories()
    generate_python_classes_from_proto()
    fix_imports()
    insert_init_files()
    setup()
 

if __name__ == "__main__":
    run()