daml/dev-env/bin/da-ghci

#!/usr/bin/env python3
#
# Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Bazel-aware wrapper for GHCi.

This script invokes GHCi loading a Haskell Bazel target and all its local
dependencies by source. If a source file is specified, then its Bazel target
will be determined automatically.

All externall Haskell dependencies and C library dependencies will be loaded
as binaries.

The REPL will be executed in the repository root.

"""

import argparse
import os
import re
import signal
import subprocess
import sys
import tempfile

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
            "--data",
            action="store",
            dest="data",
            choices=["auto", "yes","no"],
            default="auto",
            help=(
                "Whether to load data dependencies into the REPL. "
                "'yes' can imply that additional targets need to be built. "
                "If that fails you may want to retry with 'no'. "
                "'auto' will try 'yes' first, and fall back to 'no' on failure."))
    parser.add_argument(
            "target",
            metavar="TARGET",
            nargs="?",
            default="//:repl",
            help="Bazel target or source file to load into GHCi. Default: //:repl")
    parser.add_argument(
            "ghci_args",
            metavar="...",
            nargs=argparse.REMAINDER,
            help="Forward remaining arguments to GHCi.")
    args = parser.parse_args()

    if os.path.isfile(args.target):
        # Loading a source file.
        source_file = args.target

        # Discover the source file's module name.
        module_name = parse_module_name(source_file)

        # Discover the Haskell target.
        target = query_haskell_target(source_file)

        # This isn't a haskell_repl target
        is_repl = False
    else:
        # No specific module to load.
        module_name = None

        # Assume we were given a Bazel target.
        target = args.target

        # Find out if it's a haskell_repl target
        is_repl = 'haskell_repl' in query_target_kind(target)

    # In auto mode we try to build with data "yes" first, and fall back to data
    # "no" on failure. We separate the build step so that we can check for
    # build failure in isolation. In non-auto mode we don't need to separate
    # the build step and can call `bazel run` right away.
    with_data = args.data == "yes" or args.data == "auto"
    bazel_args = mk_bazel_args(target, with_data=with_data, is_repl=is_repl)
    if args.data == "auto":
        try:
            bazel_build(bazel_args)
        except subprocess.CalledProcessError:
            print("WARNING: Build with runfiles failed. Retrying without runfiles.", file=sys.stderr)
            bazel_args = mk_bazel_args(target, with_data=False, is_repl=is_repl)
            try:
                bazel_build(bazel_args)
            except subprocess.CalledProcessError:
                sys.exit(1)

    try:
        run_repl(bazel_args, args.ghci_args, module_name)
    except subprocess.CalledProcessError:
        sys.exit(1)


def bazel_build(bazel_args):
    subprocess.run(["bazel", "build"] + bazel_args, check=True)


def run_repl(bazel_args, ghci_args, module_name):
    try:
        if module_name:
            # Generate a -ghci-script that loads the module.
            script_fd, script_path = tempfile.mkstemp(text=True)
            os.write(script_fd, ":m {}".format(module_name).encode())
            ghci_args = ["-ghci-script", script_path] + ghci_args

        # Ignore SIGINT, so that Ctrl-C in GHCi doesn't kill the outer process.
        signal.signal(signal.SIGINT, signal.SIG_IGN)

        # Start GHCi.
        subprocess.run(["bazel", "run"] + bazel_args + ["--"] + ghci_args, check=True)
    finally:
        if module_name:
            os.close(script_fd)
            os.remove(script_path)


def mk_bazel_args(target, with_data, is_repl):
    bazel_args = [target if is_repl else "{}@ghci".format(target)]

    if with_data:
        bazel_args += ["--define", "ghci_data=True"]

    return bazel_args

def query_target_kind(target):
    # See query_haskell_target below.
    subprocess.run(
            ["bazel", "fetch", "--keep_going", "//..."],
            stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL)

    ret = subprocess.run(
        ["bazel", "query", "--output", "label_kind", target],
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL
    )
    if ret.returncode != 0:
        print("ERROR: Could not determine target kind.", file=sys.stderr)
        sys.exit(1)
    return ret.stdout.decode()

def query_haskell_target(source_file):
    # If Bazel needs to reload external dependencies, e.g. on a fresh checkout,
    # this could trigger additional output on stdout, e.g. due to Nix, or a
    # configure script. Here, we explicitly fetch all external dependencies, so
    # that no output from fetching leaks into the bazel query output below.
    # Bazel does not seem to have a dedicated flag to prevent such output. If
    # no fetch is necessary this call costs around a second.
    subprocess.run(
            ["bazel", "fetch", "--keep_going", "//..."],
            stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL)

    query = 'kind("haskell_*", allrdeps({}, 1))'.format(source_file)
    args = [
            "bazel", "query",
            "--universe_scope=//...", "--order_output=no",
            "--output", "label",
            "--keep_going",
            query,
            ]

    ret = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    targets = ret.stdout.decode().splitlines()

    if ret.returncode != 0 or len(targets) == 0:
        print("ERROR: No Haskell target found for the given source file.", file=sys.stderr)
        sys.exit(1)
    elif len(targets) > 1:
        print("WARNING: Found more than one Haskell target:", file=sys.stderr)
        print("\n".join("  " + t for t in targets))
        print("Loading {}".format(targets[0]))

    return targets[0]


def parse_module_name(source_file):
    prog = re.compile(".*^\s*module\s+([a-zA-Z._']+)", re.DOTALL|re.MULTILINE)
    with open(source_file, "r") as f:
        match = prog.match(f.read())
    if match:
        return match.group(1)
    else:
        print("WARNING: Could not identify module name. Will just load target.", file=sys.stderr)
        return None


if __name__ == "__main__":
    main()