codegen.py

#!/usr/bin/env python3
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Industrial-grade Protobuf/gRPC Code Generator for python.
"""

import sys
import re
import shutil
import logging
import argparse
from pathlib import Path
from typing import List, Optional

# Configure Logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - [%(levelname)s] - %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

try:
    from grpc_tools import protoc
except ImportError:
    logger.critical("Error: 'grpcio-tools' is not installed.")
    logger.critical("Please install it via: pip install grpcio-tools")
    sys.exit(1)


class CodeGenerator:
    def __init__(self, proto_root: Path, output_dir: Path, proto_files: List[str]):
        self.proto_root = proto_root.resolve()
        self.output_dir = output_dir.resolve()
        self.proto_files = proto_files
        self._check_paths()

    def _check_paths(self):
        """Validates input paths."""
        if not self.proto_root.exists():
            raise FileNotFoundError(
                f"Proto root directory not found: {self.proto_root}"
            )

        for p_file in self.proto_files:
            if not (self.proto_root / p_file).exists():
                raise FileNotFoundError(
                    f"Proto file not found: {self.proto_root / p_file}"
                )

    def _prepare_output_dir(self):
        """Creates or cleans the output directory and adds __init__.py."""
        if not self.output_dir.exists():
            self.output_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"Created output directory: {self.output_dir}")

        init_file = self.output_dir / "__init__.py"
        if not init_file.exists():
            init_file.write_text(
                "# Generated by scripts/codegen.py\n"
                "# This package contains generated protobuf code.\n"
                "# Do not edit these files manually.\n",
                encoding="utf-8",
            )
            logger.info(f"Created package marker: {init_file}")

    def _get_protoc_args(self) -> List[str]:
        """Constructs the protoc command arguments."""
        args = [
            "grpc_tools.protoc",
            f"-I{self.proto_root}",
            f"--python_out={self.output_dir}",
            f"--grpc_python_out={self.output_dir}",
        ]

        plugin_path = self._resolve_protoc_gen_mypy()
        if plugin_path is not None:
            args.append(f"--plugin=protoc-gen-mypy={plugin_path}")
            args.extend(
                [
                    f"--mypy_out={self.output_dir}",
                    f"--mypy_grpc_out={self.output_dir}",
                ]
            )
            logger.info("Enabled mypy-protobuf type generation.")
        else:
            logger.warning(
                "mypy-protobuf not found. Skipping .pyi generation."
            )

        for p_file in self.proto_files:
            args.append(str(self.proto_root / p_file))

        return args

    @staticmethod
    def _resolve_protoc_gen_mypy() -> Optional[str]:
        """
        Resolves the protoc-gen-mypy executable path.
        Protoc looks for plugins via PATH; when run from make/CI, venv bin may
        not be on PATH. Prefer the executable next to sys.executable,
        then PATH.
        """
        # 1. Same dir as current Python (venv bin when run via make)
        bin_dir = Path(sys.executable).resolve().parent
        for name in ("protoc-gen-mypy", "protoc-gen-mypy.exe"):
            candidate = bin_dir / name
            if candidate.exists() and (
                candidate.is_file() or candidate.is_symlink()
            ):
                return str(candidate)

        # 2. On PATH
        which = shutil.which("protoc-gen-mypy")
        if which:
            return which

        return None

    def _fix_imports(self):
        """
        Fixes the relative import issue in generated gRPC files.
        Converts 'import xxx_pb2' to 'from . import xxx_pb2'.
        """
        logger.info("Scanning for imports to fix...")

        import_pattern = re.compile(r"^import (\w+_pb2)(.*)$", re.MULTILINE)
        fixed_count = 0

        for py_file in self.output_dir.glob("*_pb2_grpc.py"):
            text = py_file.read_text(encoding="utf-8")
            new_text, n = import_pattern.subn(r"from . import \1\2", text)

            if n > 0:
                py_file.write_text(new_text, encoding="utf-8")
                logger.debug(f"Fixed {n} imports in {py_file.name}")
                fixed_count += 1

        if fixed_count > 0:
            logger.info(f"Successfully fixed imports in {fixed_count} files.")
        else:
            logger.info("No files needed import fixes.")

    def run(self):
        """Orchestrates the generation process."""
        try:
            self._prepare_output_dir()

            args = self._get_protoc_args()
            logger.info(f"Running protoc for: {self.proto_files}")

            exit_code = protoc.main([""] + args[1:])
            if exit_code != 0:
                logger.error(f"Protoc failed with exit code {exit_code}")
                sys.exit(exit_code)

            logger.info("Protoc compilation successful.")
            self._fix_imports()
            logger.info(f"Code generation complete. Output: {self.output_dir}")

        except Exception:
            logger.exception(
                "An unexpected error occurred during code generation."
            )
            sys.exit(1)


def main():
    current_script = Path(__file__).resolve()
    project_root = current_script.parent.parent
    default_proto_root = project_root.parent.parent / "protocol" / "proto"
    default_out_dir = project_root / "src" / "fs_client" / "_proto"

    parser = argparse.ArgumentParser(
        description="Generate python code from Proto files."
    )
    parser.add_argument(
        "--proto-root",
        type=Path,
        default=default_proto_root,
        help="Root directory containing .proto files",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=default_out_dir,
        help="Output directory for generated python code",
    )
    parser.add_argument(
        "files",
        nargs="*",
        default=["function_stream.proto"],
        help=(
            "Specific .proto files to compile "
            "(default: function_stream.proto)"
        ),
    )

    args = parser.parse_args()

    logger.info("Starting Code Generation...")
    logger.info(f"  Proto Root: {args.proto_root}")
    logger.info(f"  Output Dir: {args.out_dir}")

    generator = CodeGenerator(
        proto_root=args.proto_root,
        output_dir=args.out_dir,
        proto_files=args.files,
    )
    generator.run()


if __name__ == "__main__":
    main()