TP3/shell/tests/calc_cli.py

#!/usr/bin/env python3
"""Launch CALC.COM under the project's qemu boot harness.

Examples:
    python3 tests/calc_cli.py
    python3 tests/calc_cli.py --keys '\\r\\r/Q'
    python3 tests/calc_cli.py --keys-file keys.txt --raw > screen.bin
    python3 tests/calc_cli.py --com /path/to/CALC.COM --keys '\\r/Q'

The default keys are the welcome-return, a second return for the main loop and
`/Q' in CALC's command menu, so the command terminates instead of waiting for a
keyboard.  `--keys' accepts Python-style escapes (\\r, \\n, \\x1b, ...).
"""
import argparse
import codecs
import os
import shutil
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
SHELL = os.path.dirname(HERE)
RES = os.path.join(os.path.dirname(SHELL), "Resources", "TP3")
sys.path.insert(0, HERE)
import run_com_exec as r  # noqa: E402


def build_calc_com(verbose=False):
    """Compile Resources/TP3/CALC.PAS in a temporary include directory."""
    td = tempfile.mkdtemp(prefix="tpcalc-build-")
    for name in ("CALC.PAS", "CALC.INC"):
        shutil.copy(os.path.join(RES, name), os.path.join(td, name))
    exe = r.ensure_comtest()
    src = os.path.join(td, "CALC.PAS")
    p = subprocess.run([exe], cwd=td, input=(src + "\n").encode(),
                       capture_output=True)
    coms = [os.path.join(td, n) for n in os.listdir(td)
            if n.upper().endswith(".COM")]
    if p.returncode != 0 or not coms:
        raise RuntimeError("comtest did not produce CALC.COM:\n%s"
                           % p.stdout.decode(errors="replace"))
    com = open(coms[0], "rb").read()
    if verbose:
        sys.stderr.write(p.stdout.decode(errors="replace"))
        sys.stderr.write("built %d-byte CALC.COM\n" % len(com))
    return com, td


def decode_keys(text):
    """Decode the small escape language accepted by --keys."""
    return codecs.decode(text, "unicode_escape").encode("latin-1")


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--com", help="use an existing CALC.COM instead of building")
    ap.add_argument("--keys", default="\\r\\r/Q",
                    help="keys to feed CALC (Python escapes; default: welcome+quit)")
    ap.add_argument("--keys-file", help="read raw key bytes from this file")
    ap.add_argument("--timeout", type=int, default=15,
                    help="qemu timeout in seconds (default 15)")
    ap.add_argument("--raw", action="store_true",
                    help="write CALC's serial output as raw bytes")
    ap.add_argument("--quiet", action="store_true",
                    help="only print the exit status")
    ap.add_argument("--verbose", action="store_true",
                    help="show the compiler output and image size")
    args = ap.parse_args(argv)

    td = None
    try:
        if args.com:
            com = open(args.com, "rb").read()
        else:
            com, td = build_calc_com(args.verbose)
        if len(com) > r.COM_MAX:
            sys.exit("CALC.COM is %d bytes; loader maximum is %d"
                     % (len(com), r.COM_MAX))
        keys = (open(args.keys_file, "rb").read()
                if args.keys_file else decode_keys(args.keys))
        boot = r.build_boot_sector()
        with tempfile.TemporaryDirectory(prefix="tpcalc-run-") as runtd:
            img = os.path.join(runtd, "calc.img")
            open(img, "wb").write(r.build_floppy(boot, com, keys))
            old_timeout = r.TIMEOUT
            r.TIMEOUT = args.timeout
            try:
                out, rc, note = r.run_qemu(img)
            finally:
                r.TIMEOUT = old_timeout
        code = r.exit_code_of(rc)
        if args.raw:
            sys.stdout.buffer.write(out)
        elif not args.quiet:
            if note:
                sys.stderr.write(note + "\n")
            sys.stderr.write("CALC.COM: %d bytes, exit=%s\n" % (len(com), code))
            sys.stdout.write(r.visible_output(out))
            if out and not out.endswith(b"\n"):
                sys.stdout.write("\n")
        return 0 if code == 0 else 1
    finally:
        if td:
            shutil.rmtree(td, ignore_errors=True)


if __name__ == "__main__":
    sys.exit(main())