oiio_wheels_repair example

README.md

Example of using custom CIBW repair-wheel-command to invoke "stub_generation.py" (not included, assumed to be in the same directory as the "oiio_wheels_repair.py" file

Here's what happens:

1. Just before the build initializes, uv is installed to the local CIBW instance.
2. The wheel is built as normal...
3. We use uv run --no-config to run our custom script: A. The cross-platform repairwheel utility produces a "repaired" wheel consistent with what CIBW normally does if you don't specify a custom repair-wheel-command B. (Optionally, we strip out stuff we know we don't want -- this isn't doing much now, but if/when we start bundling dynamic (e.g. LGPL) libraries with the distribution, we'll want to strip out redundant "namelink copies" of shared libs to reduce wheel file size) C. Invoke the stub_generator.py script with uv run --with <PATH TO REPAIRED WHEEL> --no-config stub_generator.py <OUTPUT_DIR> D. Patch the generated OpenImageIO/init.pyi back into the repaired wheel, renamed as OpenImageIO/OpenImageIO.pyi

oiio_wheels_repair.py

#!/usr/bin/env python
# Copyright Contributors to the OpenImageIO project.
# SPDX-License-Identifier: Apache-2.0
# https://github.com/AcademySoftwareFoundation/OpenImageIO

# /// script
# dependencies = [
#     "repairwheel",
#     "uv",
# ]
# ///

from logging import getLogger, StreamHandler, Formatter
import os
from pathlib import Path
import shutil
import subprocess
from tempfile import TemporaryDirectory
from typing import Union, Optional, List, NoReturn

from repairwheel._vendor.auditwheel.wheeltools import InWheelCtx

logger = getLogger(__name__)
logger.addHandler(StreamHandler())
logger.setLevel(10)
logger.setFormatter(Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
logger.propagate = False

DIRS_TO_REMOVE = ["lib", "lib64"]


def wheel_slim(
    wheel_path: Union[str, Path],
    dirs_to_remove: Optional[List[Union[str, Path]]] = DIRS_TO_REMOVE,
) -> Path:
    """
    Edit the contents of the repaired wheel to remove the 'lib' directories.
    """
    with InWheelCtx(wheel_path, wheel_path) as ctx:
        root = Path(ctx.path)
        for dir_name in dirs_to_remove:
            this_dir = root / "OpenImageIO" / dir_name
            if this_dir.exists():
                shutil.rmtree(this_dir)

    return Path(wheel_path)


def wheel_stubgen(wheel_path: Union[str, Path]):
    """
    Generate stubs for the wheel at `wheel_path` and patch them into the wheel.
    """
    # Assumes stub_generator.py lives next to this script.
    _stubgen_src = Path(__file__).parent / "stub_generator.py"
    with TemporaryDirectory() as temp_dir:
        subprocess.run(["uv", "run",  "--no-config",  f"--with {wheel_path}", str(_stubgen_src), str(temp_dir)], check=True)
        with InWheelCtx(wheel_path, wheel_path) as ctx:
            root = Path(ctx.path) / "OpenImageIO"
            logger.info(f"Patching stubs into {root}/OpenImageIO.pyi")
            shutil.copy(Path(temp_dir)/"OpenImageIO/__init__.pyi", root/"OpenImageIO.pyi")
    return wheel_path

def wheel_repair(
    wheel_path: Union[str, Path],
    output_dir: Union[str, Path],
    build_dir: Optional[Union[str, Path]] = os.getenv("SKBUILD_BUILD_DIR"),
    slim: bool = False,
    include_sys_paths: bool = False,
    stubgen: bool = False,
) -> NoReturn:
    """
    Slim down and repair the wheel file at `wheel_path` with libraries from `build_dir` and save the result to `output_dir`.
    """

    with TemporaryDirectory() as temp_dir:
        if not build_dir:
            # If build_dir is not provided and $SKBUILD_BUILD_DIR is not set, unzip the wheel to a temporary wheel
            # and use that as the build directory.
            temporary_build_dir = TemporaryDirectory()
            subprocess.run(["unzip", "-q", "-d", temporary_build_dir.name, wheel_path], check=True)
            build_dir = Path(temporary_build_dir.name)/"OpenImageIO"

        if slim:
            wheel_path = wheel_slim(wheel_path)
        
        repair_wheel_command = [
            "repairwheel",
            "-l",
            f"{build_dir}/deps/dist/lib",
            "-l",
            f"{build_dir}/lib",
            "-o",
            temp_dir,
            "" if include_sys_paths else "--no-sys-paths",
            wheel_path,
        ]
        subprocess.run(repair_wheel_command, check=True)

        logger.info(f"Repaired + slimmed wheel created at {temp_dir}")
        patched_wheel_path = list(Path(temp_dir).glob("*.whl"))[0]

        if stubgen:
            patched_wheel_path = wheel_stubgen(patched_wheel_path)
            logger.info(f"Stubs patched into wheel!")

        shutil.copy(patched_wheel_path, output_dir)
        logger.info(f"Repaired + slimmed + stubbed wheel copied to {output_dir}/{patched_wheel_path.name}")


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Slim down and repair a wheel file.")
    
    parser.add_argument(
        "-w", "--wheel-path", type=str, required=True, help="Path to the wheel file",
    )
    parser.add_argument(
        "-b",
        "--build-dir",
        type=str,
        default=os.getenv("SKBUILD_BUILD_DIR"),
        help="Path to the build directory (or use $SKBUILD_BUILD_DIR)",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        type=str,
        required=True,
        help="Directory to save the repaired wheel",
    )
    parser.add_argument(
        "-t",
        "--trim",
        action="store_true",
        help="Remove everything but mandatory files from the wheel",
    )

    parser.add_argument(
        "-s",
        "--stubgen",
        action="store_true",
        help="Generate stubs and patch them into the wheel",
    )
    args = parser.parse_args()

    wheel_repair(
        build_dir=args.build_dir, 
        wheel_path=args.wheel_path, 
        output_dir=args.output_dir, 
        slim=args.trim,
        stubgen=args.stubgen,
    )


if __name__ == "__main__":
    main()

pyproject.toml

[tool.cibuildwheel]
before-build = "pip install uv"
repair-wheel-command = "uv run --no-config ./src/build-scripts/oiio_wheels_repair.py --wheel-path {wheel} --output-dir {dest_dir} --stubgen"  # --trim