setup.py

from pathlib import Path
from setuptools import setup
from setuptools.command.develop import develop as _develop
from setuptools.command.install import install as _install
from setuptools.command.build_py import build_py as _build_py
from grpc_tools import protoc
import importlib.resources


def install_and_compile_proto():
    proto_dir = Path(__file__).absolute().parent.joinpath("snet", "sdk", "resources", "proto")

    grpc_protos_include = str(importlib.resources.files("grpc_tools").joinpath("_proto"))

    print(f"Proto directory: {proto_dir}")
    print(f"Grpc include directory: {grpc_protos_include}")

    if not proto_dir.exists():
        print(f"Warning: Proto directory not found at {proto_dir}")
        return

    for fn in proto_dir.glob("*.proto"):
        print(f"Compiling protobuf: {fn}")

        command = [
            "grpc_tools.protoc",
            f"-I{proto_dir}",
            f"-I{grpc_protos_include}",
            f"--python_out={proto_dir}",
            f"--grpc_python_out={proto_dir}",
            str(fn),
        ]

        if protoc.main(command) != 0:
            print(f"Error: Failed to compile {fn}")
            raise RuntimeError(f"Protocol buffer compilation failed for {fn}")


class build_py(_build_py):
    def run(self):
        self.execute(install_and_compile_proto, (), msg="Compile protocol buffers")
        _build_py.run(self)


class develop(_develop):
    def run(self):
        self.execute(install_and_compile_proto, (), msg="Compile protocol buffers")
        _develop.run(self)


class install(_install):
    def run(self):
        self.execute(install_and_compile_proto, (), msg="Compile protocol buffers")
        _install.run(self)


setup(
    cmdclass={
        "develop": develop,
        "install": install,
        "build_py": build_py,
    },
)