buck2 tar rules

BUCK

# prelude-replay/tar/BUCK
export_file(
    name="tar_file.py",
    visibility=["PUBLIC"],
)

tar.bzl

# prelude-replay/tar.bzl

load("@prelude//python:toolchain.bzl", "PythonToolchainInfo")
load("//tar:toolchain.bzl", "TarToolchainInfo")


def tar_file_impl(ctx: AnalysisContext) -> list[Provider]:
    tar_output_name = (
        ctx.attrs.out if ctx.attrs.out else "{}.tar".format(ctx.label.name)
    )
    tar_file = ctx.actions.declare_output(tar_output_name)
    srcs = ctx.attrs.srcs

    srcs_file_cmd = cmd_args()
    for src in srcs:
        srcs_file_cmd.add(src)

    srcs_file_path = ctx.actions.write("srcs.txt", srcs_file_cmd)

    tar_toolchain = ctx.attrs._tar_toolchain[TarToolchainInfo]

    cmd = cmd_args(
        ctx.attrs._python_toolchain[PythonToolchainInfo].interpreter,
        tar_toolchain.tar_file[DefaultInfo].default_outputs,
    )
    cmd.add("--compress")
    cmd.add("true" if ctx.attrs.compress else "false")
    cmd.add("--file_path")
    cmd.add(srcs_file_path)
    cmd.add("--filename")
    cmd.add(tar_file.as_output())
    cmd.hidden(srcs)

    ctx.actions.run(cmd, category="tar")

    return [DefaultInfo(default_output=tar_file)]


tar_file = rule(
    impl=tar_file_impl,
    attrs={
        "compress": attrs.bool(default=False),
        "srcs": attrs.list(attrs.source()),
        "out": attrs.option(attrs.string(), default=None),
        "_python_toolchain": attrs.toolchain_dep(
            default="toolchains//:python",
            providers=[PythonToolchainInfo],
        ),
        "_tar_toolchain": attrs.toolchain_dep(
            default="toolchains//:tar",
            providers=[TarToolchainInfo],
        ),
    },
)

tar_file.py

# prelude-replay/tar/tar_file.py
import argparse
import tarfile
import os
import sys


def create_tar(paths, compress, filename):
    with tarfile.open(filename, "w:gz" if compress == "true" else "w") as tar:
        for path in paths:
            if not os.path.exists(path):
                raise FileNotFoundError(f"Path not found: {path}")
            tar.add(path, arcname=os.path.basename(path))


def read_paths(file_path):
    try:
        with open(file_path, "r") as file:
            return file.read().splitlines()
    except IOError as e:
        print(f"Failed to read file {file_path}: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Create a tar file from specified paths"
    )
    parser.add_argument(
        "--compress", required=True, help="Whether to gzip the tar file"
    )
    parser.add_argument(
        "--file_path",
        required=True,
        help="Path to the file containing paths to include in the tar file",
    )
    parser.add_argument("--filename", required=True, help="Name of the tar file")
    args = parser.parse_args()

    paths = read_paths(args.file_path)

    try:
        create_tar(paths, args.compress, args.filename)
    except FileNotFoundError as e:
        print(e, file=sys.stderr)
        sys.exit(1)

toolchain.bzl

# prelude-replay/tar/toolchain.bzl

TarToolchainInfo = provider(
    fields={
        "tar_file": typing.Any,
    }
)


def tar_toolchain_impl(ctx) -> list[[DefaultInfo, TarToolchainInfo]]:
    """
    A Tar toolchain.
    """
    return [
        DefaultInfo(),
        TarToolchainInfo(
            tar_file=ctx.attrs._tar_file,
        ),
    ]


tar_toolchain = rule(
    impl=tar_toolchain_impl,
    attrs={
        "_tar_file": attrs.dep(
            default="prelude-replay//tar:tar_file.py",
        ),
    },
    is_toolchain_rule=True,
)

toolchains.BUCK

load("@prelude-replay//tar:toolchain.bzl", "tar_toolchain")
tar_toolchain(
    name="tar",
    visibility=["PUBLIC"],
)