orisys-sdk/scripts/build_nuitka.py

import os
import sys
import shutil
import subprocess
from pathlib import Path
from importlib.machinery import EXTENSION_SUFFIXES

INCLUDE_LAYOUT = {
    "configs": {
        "source_dir": Path("config"),
        "target_subdir": Path("configs"),
        "patterns": ["*.json"],
    },
    "examples": {
        "source_dir": Path("examples"),
        "target_subdir": Path("examples"),
        "files": [
            Path("examples/basic_pipeline.py"),
            Path("examples/two_sensor_pipeline.py"),
            Path("examples/plot_data.py"),
        ],
    },
    "docs": {
        "source_dir": Path("docs"),
        "target_subdir": Path("docs"),
        "patterns": ["*.md", "*.pdf"],
    },
}


def install_nuitka():
    print("Checking Nuitka installation...")
    try:
        import nuitka  # noqa: F401
        print("Nuitka is installed.")
    except ImportError:
        print("Installing Nuitka...")
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "nuitka", "ordered-set", "zstandard"]
        )


def get_version() -> str:
    init_file = Path("src/orisys/__init__.py")
    if not init_file.exists():
        return "0.3.1"
    for line in init_file.read_text(encoding="utf-8").splitlines():
        if line.startswith("__version__"):
            return line.split("=", 1)[1].strip().strip('"').strip("'")
    return "0.3.1"


def find_compiled_extension(output_dir: Path) -> Path:
    candidates = []
    for suffix in EXTENSION_SUFFIXES:
        candidates.extend(output_dir.glob(f"*{suffix}"))
    if not candidates:
        all_files = ", ".join(p.name for p in output_dir.iterdir()) if output_dir.exists() else "(none)"
        raise RuntimeError(f"Nuitka did not produce a compiled extension. Found: {all_files}")
    # Prefer the shortest filename (usually the package extension itself)
    return sorted(candidates, key=lambda p: len(p.name))[0]


def move_wheel_to_sdk(wheel_path: Path, version: str) -> Path:
    sdk_version_dir = Path("SDK") / version
    sdk_version_dir.mkdir(parents=True, exist_ok=True)
    dest_wheel = sdk_version_dir / wheel_path.name
    if dest_wheel.exists():
        dest_wheel.unlink()
    shutil.move(str(wheel_path), str(dest_wheel))
    return dest_wheel


def copy_included_files(pkg_dir: Path) -> dict[str, int]:
    copied_counts = {}
    for group_name, group_cfg in INCLUDE_LAYOUT.items():
        src_dir = group_cfg["source_dir"]
        dst_dir = pkg_dir / group_cfg["target_subdir"]
        dst_dir.mkdir(exist_ok=True)

        copied = 0
        if "files" in group_cfg:
            for file_path in group_cfg["files"]:
                if file_path.exists():
                    shutil.copy(file_path, dst_dir)
                    copied += 1
                    print(f"Copied {group_name}: {file_path.name}")
        else:
            for pattern in group_cfg.get("patterns", []):
                for src_file in src_dir.glob(pattern):
                    shutil.copy(src_file, dst_dir)
                    copied += 1
                    print(f"Copied {group_name}: {src_file.name}")

        if group_name == "examples":
            (dst_dir / "__init__.py").write_text("# Examples for orisys package\n", encoding="utf-8")

        copied_counts[group_name] = copied
    return copied_counts


def get_package_data_entries(ext_suffix: str) -> dict[str, list[str]]:
    package_data = {"orisys": [f"*{ext_suffix}"], "orisys.examples": ["*.py"]}
    for group_name, group_cfg in INCLUDE_LAYOUT.items():
        if group_name == "examples":
            continue
        target_subdir = str(group_cfg["target_subdir"]).replace("\\", "/")
        for pattern in group_cfg.get("patterns", []):
            package_data["orisys"].append(f"{target_subdir}/{pattern}")
    return package_data


def build():
    install_nuitka()

    version = get_version()
    output_dir = Path("build_nuitka")
    dist_dir = Path("dist_nuitka")

    if output_dir.exists():
        shutil.rmtree(output_dir)
    if dist_dir.exists():
        shutil.rmtree(dist_dir)
    output_dir.mkdir()
    dist_dir.mkdir()

    print("\n=== Compiling with Nuitka ===")
    env = os.environ.copy()
    env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")

    cmd = [
        sys.executable, "-m", "nuitka",
        "--module",
        "--include-package=orisys",
        "--output-dir=" + str(output_dir),
        "--remove-output",
        "--no-pyi-file",
        "src/orisys",
    ]
    print(f"Running: {' '.join(cmd)}")
    subprocess.check_call(cmd, env=env)

    ext_file = find_compiled_extension(output_dir)
    ext_suffix = "".join(ext_file.suffixes)
    print(f"Generated compiled extension: {ext_file}")

    print("\n=== Preparing Wheel Structure ===")
    wheel_root = dist_dir / "wheel_root"
    pkg_dir = wheel_root / "orisys"
    pkg_dir.mkdir(parents=True)

    dest_ext = pkg_dir / f"__init__{ext_suffix}"
    shutil.copy(ext_file, dest_ext)
    print(f"Copied compiled package to {dest_ext}")

    copied_counts = copy_included_files(pkg_dir)
    print(f"Copied {copied_counts['configs']} JSON config files")
    print(f"Copied {copied_counts['examples']} example files to orisys/examples/")
    print(f"Copied {copied_counts['docs']} doc files to orisys/docs/")
    package_data = get_package_data_entries(ext_suffix)

    setup_py_path = dist_dir / "setup.py"
    setup_content = f"""
from setuptools import setup, Distribution

class BinaryDistribution(Distribution):
    def has_ext_modules(self):
        return True

setup(
    name="orisys",
    version="{version}",
    description="Orisys - Optical Tactile Sensor SDK (Nuitka Compiled)",
    author="Orisys",
    packages=['orisys', 'orisys.examples'],
    package_dir={{'': 'wheel_root'}},
    package_data={package_data},
    include_package_data=True,
    install_requires=[
        "numpy>=1.24.4",
        "opencv-python>=4.10.0.84",
        "scipy>=1.10.1",
    ],
    extras_require={{
        "cuda12": ["cupy-cuda12x>=12.3.0"],
        "cuda13": ["cupy-cuda13x>=13.0.0"],
        "cuda": ["cupy-cuda12x>=12.3.0"],
    }},
    distclass=BinaryDistribution,
    zip_safe=False,
)
"""
    setup_py_path.write_text(setup_content, encoding="utf-8")

    print("\n=== Building Wheel ===")
    subprocess.check_call([sys.executable, "setup.py", "bdist_wheel"], cwd=dist_dir)

    wheels = sorted((dist_dir / "dist").glob("*.whl"))
    if wheels:
        final_wheel = move_wheel_to_sdk(wheels[0], version)
        print(f"\n[SUCCESS] Wheel created and moved to: {final_wheel}")
    else:
        print("\n[ERROR] Wheel not found.")


if __name__ == "__main__":
    build()