nix-dev

#!/usr/bin/env python3

import argparse
import glob
import hashlib
import json
import os
import os.path
import subprocess

#
# General Helpers
#


def unlink_if_exists(path):
    if os.path.lexists(path):
        os.unlink(path)


def get_cache_directory():
    xdg_cache_dir = os.getenv("XDG_CACHE_DIR")
    if xdg_cache_dir:
        return xdg_cache_dir
    else:
        home = os.getenv("HOME")
        return f"{home}/.cache"


def get_my_cache_directory():
    cache_directory = get_cache_directory()
    return f"{cache_directory}/nix-dev"


#
# Config Related
#


def build_config(args):
    return {
        "attr": args.attr,
        "shell_nix": os.path.abspath(args.shell_nix_file),
        "shell_nix_hash": hash_of_file(args.shell_nix_file),
    }


def hash_of_file(filename):
    with open(filename, "rb", buffering=0) as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def hash_of_config(config):
    digest = hashlib.sha256()
    digest.update(json.dumps(config, sort_keys=True).encode("utf8"))
    return digest.hexdigest()


def dump_config(config, cfg_file):
    with open(cfg_file, "wt") as f:
        json.dump(config, f, indent=2, sort_keys=True)
        f.write("\n")


#
# Development Environment Handling
#


def collect_garbage(cache_dir):
    hashes = set()
    for filename in glob.glob("*", root_dir=cache_dir):
        path = os.path.join(cache_dir, filename)

        is_garbage = (
            len(filename) != 64 + 4
            or filename[-4:] not in (".cfg", ".drv", ".out")
            or not os.path.exists(path)
        )

        if is_garbage:
            print(f"Removing broken {path}")
            os.unlink(path)
        else:
            hashes.add(filename[:64])

    for h in hashes:
        cfg_file = os.path.join(cache_dir, f"{h}.cfg")
        drv_link = os.path.join(cache_dir, f"{h}.drv")
        out_link = os.path.join(cache_dir, f"{h}.out")

        is_garbage = True
        if os.path.exists(cfg_file):
            with open(cfg_file, "rt") as f:
                config = json.load(f)
            if os.path.exists(config["shell_nix"]):
                if hash_of_file(config["shell_nix"]) == config["shell_nix_hash"]:
                    is_garbage = False

        if is_garbage:
            for path in (cfg_file, drv_link, out_link):
                print(f"Removing stale {path}")
                unlink_if_exists(path)


def nix_instantiate(args, drv_link):
    cmd = ["nix-instantiate"]
    if args.attr:
        cmd.extend(["--attr", args.attr])
    cmd.extend(["--add-root", drv_link])
    cmd.append(args.shell_nix_file)
    print("nix-dev: Calling nix-instantiate...")
    subprocess.check_call(cmd)


def nix_build(args, drv_link, out_link):
    cmd = ["nix-build", "--out-link", out_link, drv_link]
    print("nix-dev: Calling nix-build...")
    subprocess.check_call(cmd)


def nix_shell(args, drv_link):
    cmd = ["nix-shell"]
    if args.command:
        cmd.extend(["--command", args.command])
    if args.keep:
        cmd.extend(["--keep", args.keep])
    if args.pure:
        cmd.extend(["--pure"])
    if args.run:
        cmd.extend(["--run", args.run])
    cmd.append(drv_link)
    print("nix-dev: Calling nix-shell...")
    os.execvp(cmd[0], cmd)


#
# Main Entry
#


def parse_args():
    parser = argparse.ArgumentParser()

    # nix-dev options
    group = parser.add_argument_group("nix-dev")
    group.add_argument(
        "--break-cache",
        action="store_true",
        help="Ignore existing cache entry and force instantiation of environment.",
    )
    group.add_argument(
        "--collect-garbage",
        action="store_true",
        help="Removes stale environments.",
    )
    group.add_argument(
        "shell_nix_file",
        nargs="?",
        default="shell.nix",
        help="If file is not given, defaults to shell.nix.",
    )

    # nix-instantiate options
    group = parser.add_argument_group("nix-instantiate")
    group.add_argument(
        "--attr",
        "-A",
        help="Select an attribute from the top-level Nix expression being evaluated.",
    )

    # nix-shell options
    group = parser.add_argument_group("nix-shell")
    group.add_argument(
        "--command",
        help="In the environment of the derivation, run the shell command COMMAND.",
    )
    group.add_argument(
        "--keep",
        help="When a --pure shell is started, keep the listed environment variables.",
    )
    group.add_argument(
        "--pure",
        action="store_true",
        help="If this flag is specified, the environment is almost entirely cleared before the interactive shell is started, so you get an environment that more closely corresponds to the “real” Nix build.",
    )
    group.add_argument(
        "--run",
        help="Like --command, but executes the command in a non-interactive shell.",
    )

    return parser.parse_args()


def main():
    args = parse_args()
    cache_dir = get_my_cache_directory()

    if args.collect_garbage:
        collect_garbage(cache_dir=cache_dir)
        return

    config = build_config(args=args)
    filename = hash_of_config(config=config)
    cfg_file = f"{cache_dir}/{filename}.cfg"
    drv_link = f"{cache_dir}/{filename}.drv"
    out_link = f"{cache_dir}/{filename}.out"
    state_exists = all(os.path.exists(p) for p in (cfg_file, drv_link, out_link))

    if not state_exists or args.break_cache:
        os.makedirs(cache_dir, exist_ok=True)
        unlink_if_exists(drv_link)
        unlink_if_exists(out_link)

        dump_config(config=config, cfg_file=cfg_file)
        nix_instantiate(args=args, drv_link=drv_link)
        nix_build(args=args, drv_link=drv_link, out_link=out_link)

    nix_shell(args=args, drv_link=drv_link)


if __name__ == "__main__":
    main()