Files
OrcaSlicer/scripts/generate_orca_python_stubs.py
2026-07-02 17:49:36 +08:00

213 lines
6.8 KiB
Python

#!/usr/bin/env python3
"""
Generate Python .pyi stubs for OrcaSlicer's pybind11 plugin API.
The script creates a local virtual environment, installs pybind11-stubgen,
imports the built `orca` extension module, and writes stubs to ./typings by
default. It intentionally does not update editor settings.
"""
from __future__ import annotations
import argparse
import os
import shutil
import subprocess
import sys
import sysconfig
import venv
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_VENV = REPO_ROOT / ".venv-stubgen"
DEFAULT_OUTPUT = REPO_ROOT / "typings"
DEFAULT_BUILD_DIR = REPO_ROOT / "build"
DEFAULT_CONFIG = "RelWithDebInfo"
MODULE_NAME = "orca"
def log(message: str) -> None:
print(f"[orca-stubgen] {message}")
def run(command: list[str], *, env: dict[str, str] | None = None, cwd: Path = REPO_ROOT) -> None:
log(" ".join(command))
subprocess.run(command, cwd=cwd, env=env, check=True)
def venv_python(venv_dir: Path) -> Path:
if os.name == "nt":
return venv_dir / "Scripts" / "python.exe"
return venv_dir / "bin" / "python"
def ensure_venv(venv_dir: Path) -> Path:
python = venv_python(venv_dir)
if not python.exists():
log(f"creating virtual environment: {venv_dir}")
venv.EnvBuilder(with_pip=True).create(venv_dir)
return python
def ensure_stubgen(python: Path) -> None:
probe = subprocess.run(
[str(python), "-m", "pybind11_stubgen", "--help"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
if probe.returncode == 0:
return
run([str(python), "-m", "pip", "install", "--upgrade", "pip"])
run([str(python), "-m", "pip", "install", "pybind11-stubgen"])
def module_suffixes() -> list[str]:
suffixes = [".pyd", ".so", ".dylib"]
ext_suffix = sysconfig.get_config_var("EXT_SUFFIX")
if ext_suffix:
suffixes.insert(0, str(ext_suffix))
return list(dict.fromkeys(suffixes))
def is_orca_extension(path: Path) -> bool:
if not path.is_file():
return False
if path.name == f"{MODULE_NAME}.so" or path.name == f"{MODULE_NAME}.pyd":
return True
return any(path.name.startswith(f"{MODULE_NAME}.") and path.name.endswith(suffix) for suffix in module_suffixes())
def find_module_dir(build_dir: Path, config: str) -> Path | None:
candidates = [
build_dir / "src" / "slic3r" / config,
build_dir / "src" / "slic3r",
REPO_ROOT / "build" / "src" / "slic3r" / config,
REPO_ROOT / "build" / "src" / "slic3r",
REPO_ROOT / "build" / "arm64" / "src" / "slic3r" / config,
REPO_ROOT / "build" / "arm64" / "src" / "slic3r",
]
for candidate in candidates:
if not candidate.exists():
continue
if any(is_orca_extension(path) for path in candidate.iterdir()):
return candidate
for root in dict.fromkeys([build_dir, REPO_ROOT / "build", REPO_ROOT / "build" / "arm64"]):
if not root.exists():
continue
for path in root.rglob(f"{MODULE_NAME}*"):
if is_orca_extension(path):
return path.parent
return None
def build_stubgen_module(build_dir: Path, config: str) -> None:
run([
"cmake",
"--build",
str(build_dir),
"--config",
config,
"--target",
"orca_stubgen",
"--",
])
def import_env(module_dir: Path) -> dict[str, str]:
env = os.environ.copy()
existing = env.get("PYTHONPATH")
paths = [str(module_dir)]
if existing:
paths.append(existing)
env["PYTHONPATH"] = os.pathsep.join(paths)
return env
def verify_import(python: Path, module_dir: Path) -> None:
code = (
"import orca; "
"print(orca.__file__); "
"assert hasattr(orca, 'printer_agent'), 'orca.printer_agent is missing'"
)
run([str(python), "-c", code], env=import_env(module_dir))
def clean_output(output_dir: Path) -> None:
package_dir = output_dir / MODULE_NAME
module_stub = output_dir / f"{MODULE_NAME}.pyi"
if package_dir.exists():
shutil.rmtree(package_dir)
if module_stub.exists():
module_stub.unlink()
def generate_stubs(python: Path, module_dir: Path, output_dir: Path, ignore_errors: bool) -> None:
command = [str(python), "-m", "pybind11_stubgen", MODULE_NAME, "-o", str(output_dir)]
if ignore_errors:
command.append("--ignore-all-errors")
run(command, env=import_env(module_dir))
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--venv", type=Path, default=DEFAULT_VENV, help="Virtual environment path.")
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="Directory for generated stubs.")
parser.add_argument("--build-dir", type=Path, default=DEFAULT_BUILD_DIR, help="CMake build directory.")
parser.add_argument("--config", default=DEFAULT_CONFIG, help="CMake configuration name.")
parser.add_argument("--module-dir", type=Path, help="Directory containing the built orca extension module.")
parser.add_argument("--build-missing", action="store_true", help="Build the orca_stubgen target if orca is missing.")
parser.add_argument("--clean", action="store_true", help="Remove existing generated orca stubs before writing new ones.")
parser.add_argument(
"--strict",
action="store_true",
help="Do not pass --ignore-all-errors to pybind11-stubgen.",
)
return parser.parse_args()
def main() -> int:
args = parse_args()
venv_dir = args.venv.resolve()
output_dir = args.output.resolve()
build_dir = args.build_dir.resolve()
module_dir = args.module_dir.resolve() if args.module_dir else None
python = ensure_venv(venv_dir)
ensure_stubgen(python)
if module_dir is None:
module_dir = find_module_dir(build_dir, args.config)
if module_dir is None and args.build_missing:
build_stubgen_module(build_dir, args.config)
module_dir = find_module_dir(build_dir, args.config)
if module_dir is None:
print(
"Could not find a built orca extension module.\n"
f"Expected something like: {build_dir / 'src' / 'slic3r' / args.config / 'orca.so'}\n"
"Build it with:\n"
f" cmake --build {build_dir} --config {args.config} --target orca_stubgen --\n"
"Or rerun this script with --build-missing.",
file=sys.stderr,
)
return 1
output_dir.mkdir(parents=True, exist_ok=True)
if args.clean:
clean_output(output_dir)
log(f"using module directory: {module_dir}")
verify_import(python, module_dir)
generate_stubs(python, module_dir, output_dir, ignore_errors=not args.strict)
log(f"wrote stubs to: {output_dir}")
return 0
if __name__ == "__main__":
raise SystemExit(main())