common-mk/runtime_wrapper.py

#!/usr/bin/env python3
# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Wrapper to handle stdio & env vars."""

import argparse
import contextlib
import os
from pathlib import Path
import subprocess
import sys
from typing import List, Optional


class _EnvAdd(argparse.Action):
    """Add a KEY=VAL to an env dictionary."""

    def __init__(self, option_strings, dest, **kwargs):
        if "nargs" in kwargs:
            raise ValueError(f"nargs is not supported for {dest}")
        super().__init__(option_strings, dest, nargs=1, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        if getattr(namespace, self.dest, None) is None:
            setattr(namespace, self.dest, {})
        if "=" not in values[0]:
            parser.error(
                f"{option_string} must use the form KEY=VALUE, not: {values[0]}"
            )
        key, value = values[0].split("=", 1)
        getattr(namespace, self.dest)[key] = value


class _EnvDel(argparse.Action):
    """Remove a var from an env dictionary."""

    def __init__(self, option_strings, dest, **kwargs):
        if "nargs" in kwargs:
            raise ValueError(f"nargs is not supported for {dest}")
        super().__init__(option_strings, dest, nargs=1, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        if getattr(namespace, self.dest, None) is None:
            setattr(namespace, self.dest, {})
        # Mark the variable for removal so that we can process all the args at
        # the end and make sure we handle `--set-env F=1 --unset-env F`.
        getattr(namespace, self.dest)[values[0]] = None


def get_parser() -> argparse.ArgumentParser:
    """Build the argument parser."""
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument(
        "--stdin", default=None, type=Path, help="file to use as stdin"
    )

    parser.add_argument(
        "--stdout", default=None, type=Path, help="file to use as stdout"
    )

    parser.add_argument(
        "--stderr", default=None, type=Path, help="file to use as stderr"
    )

    parser.add_argument(
        "--set-env",
        dest="extra_env",
        action=_EnvAdd,
        metavar="KEY=VAL",
        help="set variable in environment",
    )

    parser.add_argument(
        "--unset-env",
        dest="extra_env",
        action=_EnvDel,
        metavar="VAR",
        help="unset variable in environment",
    )

    parser.add_argument("cmdargs", nargs="+", help="args of command to run")

    return parser


def parse_arguments(argv: List[str]) -> argparse.Namespace:
    """Parse and validate arguments."""
    parser = get_parser()
    opts = parser.parse_args(argv)
    return opts


def main(argv: List[str]) -> Optional[int]:
    opts = parse_arguments(argv)
    print(opts.extra_env)

    in_file = None
    out_file = None
    err_file = None
    kwargs = {}

    if opts.extra_env:
        env = os.environ.copy()
        env.update(opts.extra_env)
        for key, value in opts.extra_env.items():
            if value is None:
                env.pop(key)
        kwargs["env"] = env

    with contextlib.ExitStack() as stack:
        if opts.stdin:
            in_file = stack.enter_context(open(opts.stdin, "rb"))
            kwargs["stdin"] = in_file
        if opts.stdout:
            out_file = stack.enter_context(open(opts.stdout, "wb"))
            kwargs["stdout"] = out_file
        if opts.stderr:
            err_file = stack.enter_context(open(opts.stderr, "wb"))
            kwargs["stderr"] = err_file

        subprocess.check_call(opts.cmdargs, **kwargs)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))