main.py

from __future__ import annotations

import argparse
import ast
from pathlib import Path
import time

from frontend.semantic_visitor import SemanticBuilder


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build custom IR for Python dimension checking."
    )
    parser.add_argument(
        "paths",
        nargs="*",
        default=["examples"],
        help="Python files or directories to analyze",
    )
    parser.add_argument(
        "-o",
        "--output",
        default="ir_out",
        help="Output directory for generated IR files",
    )
    return parser.parse_args()


def iter_python_files(paths: list[str]):
    for raw in paths:
        path = Path(raw)
        if path.is_dir():
            yield from sorted(path.glob("*.py"))
        elif path.suffix == ".py":
            yield path


def build_file(path: Path):
    source = path.read_text(encoding="utf-8")
    tree = ast.parse(source, filename=str(path))
    builder = SemanticBuilder(
        module_name=path.stem,
        file_path=str(path),
    )
    return builder.build(tree)


def main() -> None:
    args = parse_args()
    start = time.time()

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    for path in iter_python_files(args.paths):
        ir = build_file(path)
        output_file = output_dir / f"{path.stem}.pb"

        with open(output_file, "wb") as f:
            f.write(ir.to_proto().SerializeToString())

    print(f"Successfully generated IR in: {time.time() - start:.4f} seconds")


if __name__ == "__main__":
    main()