Example wheel patch script

README.md

Example of a script to patch wheels run during the cibuildwheel repair-wheel step

repair_wheel.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",
# ]
# ///

from logging import getLogger
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__)

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_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,
) -> NoReturn:
    """
    Slim down and repair the wheel file at `wheel_path` with libraries from `build_dir` and save the result to `output_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)
    
    command = [
        "repairwheel",
        "-l",
        f"{build_dir}/deps/dist/lib",
        "-l",
        f"{build_dir}/lib",
        "-o",
        output_dir,
        "" if include_sys_paths else "--no-sys-paths",
        wheel_path,
    ]
    subprocess.run(command, check=True)

    logger.info(f"Repaired + slimmed wheel created at {output_dir}")
    


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(
        "-s",
        "--slim",
        action="store_true",
        help="Remove everything but mandatory files from 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.slim
    )


if __name__ == "__main__":
    main()