ykcheungab
111 lines
3.3 KiB
Python
111 lines
3.3 KiB
Python
"""Nuitka compilation of the orisys package."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from importlib.machinery import EXTENSION_SUFFIXES
|
|
from pathlib import Path
|
|
|
|
_BUILD_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_BUILD_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_BUILD_ROOT))
|
|
|
|
from common.paths import repo_path
|
|
|
|
NUITKA_INCLUDE_PACKAGES = (
|
|
"orisys",
|
|
"orisys.capture",
|
|
"orisys.finger",
|
|
"orisys.stitch",
|
|
"orisys.stitch.backends",
|
|
"orisys.unroll",
|
|
"orisys.hhd",
|
|
"orisys.hhd.solvers",
|
|
)
|
|
|
|
NUITKA_NOFOLLOW_IMPORTS = (
|
|
# Keep as runtime deps (shipped via wheel install_requires / extras).
|
|
"numpy",
|
|
"scipy",
|
|
"cv2",
|
|
"numba",
|
|
"orisys_blend",
|
|
# Optional accelerators / GUIs: imported behind try/except, resolved from the
|
|
# user's environment at runtime. Listing them here is not optional - without
|
|
# it, Nuitka follows whatever happens to be installed in the *build*
|
|
# environment, which blows up compile time and makes wheel contents depend on
|
|
# ambient packages. tests/test_nuitka_nofollow.py enforces full coverage.
|
|
"cupy",
|
|
"cupyx",
|
|
"pyopencl",
|
|
"pyvkfft",
|
|
"open3d",
|
|
"PyQt5",
|
|
"pyqtgraph",
|
|
"torch",
|
|
"torchvision",
|
|
)
|
|
|
|
|
|
def install_nuitka() -> None:
|
|
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 find_compiled_extension(output_dir: Path, module_name: str) -> Path:
|
|
candidates = []
|
|
for suffix in EXTENSION_SUFFIXES:
|
|
candidates.extend(output_dir.glob(f"{module_name}*{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 for {module_name}. Found: {all_files}"
|
|
)
|
|
return sorted(candidates, key=lambda p: len(p.name))[0]
|
|
|
|
|
|
def compile_orisys(build_root: Path, env: dict[str, str]) -> Path:
|
|
module_output_dir = build_root / "orisys"
|
|
if module_output_dir.exists():
|
|
shutil.rmtree(module_output_dir)
|
|
module_output_dir.mkdir(parents=True)
|
|
|
|
cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"nuitka",
|
|
"--module",
|
|
"--output-dir=" + str(module_output_dir),
|
|
"--remove-output",
|
|
"--no-pyi-file",
|
|
repo_path("src/orisys").as_posix(),
|
|
]
|
|
for package_name in NUITKA_INCLUDE_PACKAGES:
|
|
cmd.append(f"--include-package={package_name}")
|
|
for import_name in NUITKA_NOFOLLOW_IMPORTS:
|
|
cmd.append(f"--nofollow-import-to={import_name}")
|
|
|
|
print(f"Running: {' '.join(cmd)}")
|
|
subprocess.check_call(cmd, env=env)
|
|
|
|
ext_file = find_compiled_extension(module_output_dir, "orisys")
|
|
print(f"Generated compiled extension for orisys: {ext_file}")
|
|
return ext_file
|
|
|
|
|
|
def copy_compiled_orisys(pkg_dir: Path, ext_file: Path) -> str:
|
|
ext_suffix = "".join(ext_file.suffixes)
|
|
dest_ext = pkg_dir / f"__init__{ext_suffix}"
|
|
shutil.copy(ext_file, dest_ext)
|
|
print(f"Copied compiled module orisys to {dest_ext}")
|
|
return ext_suffix
|