pip_build.py

"""Script to create (and optionally install) a `.whl` archive for Keras RS.

Usage:

1. Create a `.whl` file in `dist/`:

```
python3 pip_build.py
```

2. Also install the new package immediately after:

```
python3 pip_build.py --install
```
"""

import argparse
import datetime
import os
import pathlib
import re
import shutil

package = "keras_rs"
build_directory = "tmp_build_dir"
dist_directory = "dist"
to_copy = ["pyproject.toml", "README.md"]


def ignore_files(path: str, filenames: list[str]) -> list[str]:
    if path.endswith("testing"):
        return filenames
    return [f for f in filenames if f.endswith("_test.py")]


def update_version_for_nightly(build_path: pathlib.Path, version: str) -> str:
    """Export Version and Package Name."""
    date = datetime.datetime.now()
    version += f".dev{date:%Y%m%d%H}"
    # Update `name = "keras-rs"` with "keras-rs-nightly"
    pyproj_pth = build_path / "pyproject.toml"
    pyproj_str_before = pyproj_pth.read_text()
    pyproj_str_after = pyproj_str_before.replace(
        'name = "keras-rs"', 'name = "keras-rs-nightly"'
    )
    if pyproj_str_before == pyproj_str_after:
        raise ValueError("Package name replacement failed")
    pyproj_pth.write_text(pyproj_str_after)

    # Make sure to export the __version__ string
    version_py = build_path / package / "src" / "version.py"
    version_str_before = version_py.read_text()
    version_str_after = re.sub(
        "\n__version__ = .*\n",
        f'\n__version__ = "{version}"\n',
        version_str_before,
    )
    if version_str_before == version_str_after:
        raise ValueError("Version replacement failed")
    version_py.write_text(version_str_after)
    return version


def copy_source_to_build_directory(
    root_path: pathlib.Path, build_path: pathlib.Path
) -> None:
    # Copy sources (`keras_rs/` directory and setup files) to build directory
    shutil.copytree(
        root_path / package,
        build_path / package,
        ignore=ignore_files,
    )
    for fname in to_copy:
        shutil.copy(root_path / fname, root_path / build_directory / fname)


def build_wheel(
    build_path: pathlib.Path, dist_path: pathlib.Path, __version__: str
) -> pathlib.Path:
    # Build the package
    os.chdir(build_path)
    os.system("python3 -m build")

    # Save the dist files generated by the build process
    if not os.path.exists(dist_path):
        os.mkdir(dist_path)
    for fpath in (build_path / dist_directory).glob("*.*"):
        shutil.copy(fpath, dist_path)

    # Find the .whl file path
    for fname in os.listdir(dist_path):
        if __version__ in fname and fname.endswith(".whl"):
            whl_path = dist_path / fname
            print(f"Build successful. Wheel file available at {whl_path}")
            return whl_path
    raise FileNotFoundError("Build failed")


def build(root_path: pathlib.Path, is_nightly: bool) -> pathlib.Path:
    try:
        build_path = root_path / build_directory
        dist_path = root_path / dist_directory

        if os.path.exists(build_path):
            raise ValueError(f"Directory already exists: {build_path}")
        os.mkdir(build_path)

        from keras_rs.src import version

        keras_rs_version = version.__version__
        copy_source_to_build_directory(root_path, build_path)
        if is_nightly:
            keras_rs_version = update_version_for_nightly(
                build_path, keras_rs_version
            )
        return build_wheel(build_path, dist_path, keras_rs_version)
    finally:
        # Clean up: remove the build directory (no longer needed)
        os.chdir(root_path)
        shutil.rmtree(build_path)


def install_whl(whl_fpath: pathlib.Path) -> None:
    print(f"Installing wheel file: {whl_fpath}")
    os.system(f"pip3 install {whl_fpath} --force-reinstall --no-dependencies")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--install",
        action="store_true",
        help="Whether to install the generated wheel file.",
    )
    parser.add_argument(
        "--nightly",
        action="store_true",
        help="Whether to generate nightly wheel file.",
    )
    args = parser.parse_args()
    root_path = pathlib.Path(__file__).parent.resolve()
    whl_path = build(root_path, args.nightly)
    if whl_path and args.install:
        install_whl(whl_path)