Python script for building a wheel with optional data copy-in and platform wheel tagging

build_wheel.py

import argparse
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="store_true", help="verbose")

    parser.add_argument(
        "--package-dir",
        type=pathlib.Path,
        required=True,
        help="Root directory of package",
    )

    parser.add_argument(
        "--platform-tag", help="tag for a platform wheel (e.g. macosx_12_6)"
    )

    parser.add_argument(
        "--dynamic-data-dir",
        type=pathlib.Path,
        help="Relative path from package root to dynamic data directory",
    )

    parser.add_argument(
        "--dynamic-data-file",
        type=pathlib.Path,
        action="append",
        help="Path to dynamic data file to copy to dynamic data dir",
    )

    parser.add_argument(
        "--output", type=pathlib.Path, required=True, help="Output file"
    )

    args = parser.parse_args()

    if not args.verbose:
        args.verbose = os.environ.get("VERBOSE", "0").lower() not in (
            "0",
            "false",
            "off",
        )

    if not args.package_dir.exists():
        raise FileNotFoundError(args.package_dir)

    if args.dynamic_data_file:
        if not args.dynamic_data_dir:
            msg = "--dynamic-data-dir is required with --dynamic-data-file"
            raise ValueError(msg)

        absent_files = [f for f in args.dynamic_data_file if not f.exists()]
        if absent_files:
            msg = f"Dynamic data files {absent_files} not found"
            raise ValueError(msg)

    return args


def _build_wheel(
    package_dir,
    platform_tag,
    dynamic_data_dir,
    dynamic_data_files,
    output_file,
    verbose,
):
    """Hermetically builds a wheel file and copies it to output_file."""

    # Because of race conditions, copy the package tree to a private staging dir
    with tempfile.TemporaryDirectory() as temp_dir:
        staging_dir = pathlib.Path(temp_dir) / package_dir.name
        if verbose:
            print(f"  staging-dir: {staging_dir}")

        # Copy the package source tree.
        ignore_patterns = shutil.ignore_patterns("*.pyc", "*.egg-info", "build", "dist")
        shutil.copytree(package_dir, staging_dir, ignore=ignore_patterns)

        # Dynamic data is anything from outside the package root, to be copied in.
        if dynamic_data_files:
            dynamic_data_abs = staging_dir / dynamic_data_dir
            dynamic_data_abs.mkdir(parents=True)
            if verbose:
                print(f"  dynamic data dir: {dynamic_data_abs}")

            for src_file in dynamic_data_files:
                dst_file = dynamic_data_abs / src_file.name
                if verbose:
                    print(f"  copying: {src_file} => {dst_file}")
                shutil.copyfile(src_file, dst_file, follow_symlinks=True)

        # Run the packaging command to create the .whl file in the 'dist' subdir
        build_cmd = [sys.executable, "-m", "build", "-w", "-n"]

        if verbose:
            print(f'  running: "{" ".join(build_cmd)}"')

        try:
            result = subprocess.run(
                build_cmd,
                check=True,
                cwd=staging_dir,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
            )
        except subprocess.CalledProcessError as err:
            print(err.stdout.decode())
            return err.returncode

        if verbose:
            print(result.stdout.decode())

        # The "any" wheel is built and sitting in 'dist'
        built_wheel = next(staging_dir.glob("dist/*.whl"))

        if platform_tag:
            # If the wheel is artificially specified to be a platform wheel, run
            # "wheel tags --platform-tag XYZ" to update the "any" tag to the platform tag.
            # This creates a new .whl file.

            tag_cmd = [
                sys.executable,
                "-m",
                "wheel",
                "tags",
                "--platform-tag",
                platform_tag,
                built_wheel,
            ]

            if verbose:
                print(f'  running: "{" ".join(str(a) for a in tag_cmd)}"')

            try:
                result = subprocess.run(
                    tag_cmd, check=True, cwd=staging_dir, capture_output=True
                )
            except subprocess.CalledProcessError as err:
                print(err.stdout.decode())
                return err.returncode

            if verbose:
                print(result.stdout.decode())

            # Update the wheel so the platform wheel gets copied out of the sandbox.
            built_wheel = next(staging_dir.glob(f"dist/*{platform_tag}.whl"))

        if verbose:
            print(f"  wheel: {built_wheel}")
        shutil.copyfile(built_wheel, output_file)

    return result.returncode


def main():
    args = _parse_args()
    if args.verbose:
        print(f"{pathlib.Path(__file__).stem}:")
        for arg in vars(args):
            print(f"  {arg}: {getattr(args, arg)}")

    return _build_wheel(
        args.package_dir,
        args.platform_tag,
        args.dynamic_data_dir,
        args.dynamic_data_file,
        args.output,
        args.verbose,
    )


if __name__ == "__main__":
    sys.exit(main())