Find sources that went into derivations (recursively)

query-build.py

from pynixutil import drvparse
from functools import wraps
from typing import (
    Optional,
    Callable,
    List,
    Set,
)
import subprocess
import argparse


parser = argparse.ArgumentParser(description='Find sources/build closures recursively for derivations')

subparsers = parser.add_subparsers(dest="subcommand")

sources_parser = subparsers.add_parser('find_sources')
sources_parser.add_argument('drvs', type=str, nargs='+', help='Derivation store paths')

build_closure_parser = subparsers.add_parser('build_closure')
build_closure_parser.add_argument('drvs', type=str, nargs='+', help='Derivation store paths')


def _handle_drvs_once(f: Callable):
    """Don't double-handle already handled drvs"""
    handled_drvs: Set[str] = set()

    @wraps(f)
    def handler(drv_path: str, *args, **kwargs):
        if drv_path in handled_drvs:
            return
        handled_drvs.add(drv_path)
        return f(drv_path, *args, **kwargs)

    return handler


def find_sources(drvs: List[str]) -> List[str]:
    """Find sources that went into derivation (recursively)"""
    store_paths: Set[str] = set()

    @_handle_drvs_once
    def handle_drv(drv_path) -> None:
        with open(drv_path) as f:
            drv = drvparse(f.read())

        for _, drv_output in drv.outputs.items():
            if drv_output.hash_algo is not None:  # Fixed output (sources)
                store_paths.add(drv_output.path)

        for input_src in drv.input_srcs:
            store_paths.add(input_src)

        for input_drv_path, input_drv_outputs in drv.input_drvs.items():
            handle_drv(input_drv_path)

    for drv_path in drvs:
        handle_drv(drv_path)

    return sorted(store_paths)


def find_direct_build_requisites(drvs: List[str]) -> List[str]:
    """Find the runtime closure of a _build_ based on derivation paths"""

    store_paths: Set[str] = set()

    @_handle_drvs_once
    def extract_outputs(drv_path, outputs: List[str]):
        with open(drv_path) as f:
            drv = drvparse(f.read())

        for output in outputs:
            store_paths.add(drv.outputs[output].path)

    @_handle_drvs_once
    def handle_drv(drv_path) -> None:
        with open(drv_path) as f:
            drv = drvparse(f.read())

        for input_drv, outputs in drv.input_drvs.items():
            extract_outputs(input_drv, outputs)

    for drv_path in drvs:
        handle_drv(drv_path)

    o = subprocess.check_output(["nix-store", "--query", "--requisites"] + sorted(store_paths))
    return [l for l in o.decode().split("\n") if l]


def main():
    args = parser.parse_args()

    if args.subcommand == "build_closure":
        for p in find_direct_build_requisites(args.drvs):
            print(p)

    elif args.subcommand == "find_sources":
        for p in find_sources(args.drvs):
            print(p)

    else:
        raise ValueError(f"Unhandled subcommand: {args.subcommand}")


if __name__ == "__main__":
    main()