diff --git a/bazel/labgrid/runner/BUILD.bazel b/bazel/labgrid/runner/BUILD.bazel new file mode 100644 index 0000000000000000000000000000000000000000..5c0e712128bb35a0d1603121fd5bd95e49834190 --- /dev/null +++ b/bazel/labgrid/runner/BUILD.bazel @@ -0,0 +1,35 @@ +load("@rules_labgrid//labgrid/config:transition.bzl", "labgrid_config_transition") +load("@rules_python//python:defs.bzl", "py_library") + +TOOLS = [ + "mktemp", + "mkdir", + "rm", + "env", + "mv", +] + +[ + labgrid_config_transition( + name = tool, + srcs = ["@ape//ape:{}".format(tool)], + ) + for tool in TOOLS +] + +py_library( + name = "runner", + srcs = [ + "__init__.py", + "runner.py", + ], + data = [":{}".format(tool) for tool in TOOLS], + visibility = ["//visibility:public"], + deps = [ + "//bazel/labgrid/strategy", + "//bazel/labgrid/util:target", + "//bazel/python/runfiles", + "//labgrid:pkg", + "//labgrid/config:deps", + ], +) diff --git a/bazel/labgrid/runner/__init__.py b/bazel/labgrid/runner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..31f77f5be19856314a985ee39dbacf40112306d4 --- /dev/null +++ b/bazel/labgrid/runner/__init__.py @@ -0,0 +1,3 @@ +from .runner import FileTransfer, runner + +__all__ = ["FileTransfer", "runner"] diff --git a/bazel/labgrid/runner/runner.py b/bazel/labgrid/runner/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..a2f88d36ce822cf23c97d61cc39a56c09041d4db --- /dev/null +++ b/bazel/labgrid/runner/runner.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from os import environ, linesep +from pathlib import Path, PurePath +from sys import stderr, stdout +from typing import ( + Iterator, + Mapping, +) + +from bazel.labgrid.strategy import State, transition +from bazel.labgrid.util.target import TemporaryDirectory +from bazel.python.runfiles import runfile +from labgrid import Environment +from labgrid.driver.exception import ExecutionError +from labgrid.logging import StepLogger, basicConfig + + +@dataclass +class _Tools: + mktemp: Path = runfile("ape/ape/assimilate/mktemp/mktemp") + rm: Path = runfile("ape/ape/assimilate/rm/rm") + mkdir: Path = runfile("ape/ape/assimilate/mkdir/mkdir") + env: Path = runfile("ape/ape/assimilate/env/env") + mv: Path = runfile("ape/ape/assimilate/mv/mv") + + +@dataclass(frozen=True) +class FileTransfer: + remote: PurePath + local: PurePath + optional: bool + # FIXME: Decide whether or not to transfer runfiles implicitly instead of using a flag + include_runfiles: bool + + def __init__(self, remote: str, local: str, optional=False, include_runfiles=False): + object.__setattr__(self, "remote", PurePath(remote)) + object.__setattr__(self, "local", PurePath(local)) + object.__setattr__(self, "optional", optional) + object.__setattr__(self, "include_runfiles", include_runfiles) + + +def _log_level(value): + # The Logging module converts strings to level integers with this (unintuitively named) function + log_level = logging.getLevelName(value) + if isinstance(log_level, str): + # getLevelName will return a string if the level was not recognised + raise ValueError(f"Unsupported value for LOG_LEVEL: '{value}'") + + return log_level + + +@contextmanager +def runner(): + config = PurePath(environ["LG_ENV"]) + log_level = _log_level(environ.get("LOG_LEVEL", "ERROR")) + tools = _Tools() + + runfiles_dir = None + # Allow any Bazel binaries used in the configuration to work + try: + # FIXME: Convert into context manager + runfiles_dir = environ["RUNFILES_DIR"] + del environ["RUNFILES_DIR"] + except KeyError: + pass + + # Activating the StepLogger requires a DEBUG level of verbosity + basicConfig(level=log_level) + if log_level <= logging.DEBUG: + StepLogger.start() + + target = Environment(str(config)).get_target() + strategy = target.get_driver("Strategy") + with transition(strategy, State.from_env()): + # Retrieve the communication protocols + shell = target.get_driver("CommandProtocol") + transfer = target.get_driver("FileTransferProtocol") + with TemporaryDirectory( + shell, + transfer, + tools.mktemp, + tools.rm, + tools.mkdir, + tools.mv, + ) as temp: + r = Runner(temp, transfer, shell, runfiles_dir) + + # Transfer 'env' binary over to the device + r.put([FileTransfer(r.exec_path / tools.env.name, tools.env)]) + + yield r + + +class Runner: + def __init__(self, exec_path, transfer, shell, runfiles_dir): + self.exec_path = exec_path + self.transfer = transfer + self.shell = shell + self.runfiles_dir = runfiles_dir + + def put(self, uploads: Iterator[FileTransfer]): + for upload in uploads: + remote, local = self._resolve(upload) + self.exec_path.mkdir(remote.parent) + self.transfer.put(local, remote) + + if upload.include_runfiles: + src_runfiles = local.with_suffix(".runfiles") + dest_runfiles = remote.with_suffix(".runfiles") + if Path(src_runfiles).exists(): + self.transfer.put(f"{src_runfiles}", f"{dest_runfiles}") + elif self.runfiles_dir is not None: + self.transfer.put(f"{self.runfiles_dir}", f"{dest_runfiles}") + + def get(self, downloads: Iterator[FileTransfer], code: int): + for download in downloads: + remote, local = self._resolve(download) + # FIXME: Check if file exists with `@ape//ape:test` instead of assuming any error is due to a missing file + try: + self.transfer.get(remote, local) + except ExecutionError as e: + if not download.optional and code == 0: + raise e + + def run(self, cmd: str, env: Mapping[str, str]) -> int: + # Construct the CLI using env, the env vars and then the program and args. + name_values = " ".join(f"{k}={v}" for k, v in env.items()) + cmd = f"cd {self.exec_path} && {self.exec_path / 'env'} {name_values} {cmd}" + + # Run the transferred program + out, err, code = self.shell.run(cmd) + for line in out: + stdout.write(f"{line}{linesep}") + for line in err: + stderr.write(f"{line}{linesep}") + + return code + + def _resolve(self, transfer: FileTransfer) -> tuple[str, str]: + remote = self._join(self.exec_path, transfer.remote) + local = self._join(Path.cwd(), transfer.local) + return (remote, local) + + @staticmethod + def _join(base: PurePath, path: PurePath) -> PurePath: + if not path.is_absolute(): + return base / path + return path diff --git a/bazel/labgrid/strategy/__init__.py b/bazel/labgrid/strategy/__init__.py index b1f21abb237c80df9e852ffc936c1c497731ce2d..7994b5dd3b31fc94c78c652ee654b5cdcb2fcebf 100644 --- a/bazel/labgrid/strategy/__init__.py +++ b/bazel/labgrid/strategy/__init__.py @@ -1,6 +1,6 @@ from .localhost import LocalhostStrategy from .qemu import QEMUStrategy from .ssh import SSHStrategy -from .transition import transition +from .transition import transition, State -__all__ = ["LocalhostStrategy", "QEMUStrategy", "SSHStrategy", "transition"] +__all__ = ["LocalhostStrategy", "QEMUStrategy", "SSHStrategy", "transition", "State"] diff --git a/bazel/labgrid/strategy/transition.py b/bazel/labgrid/strategy/transition.py index c7d3bf0823a446fa10a0b370b534011d4127f9bb..e81531833b10c6819c42dfcaff657f5f40ab7b09 100644 --- a/bazel/labgrid/strategy/transition.py +++ b/bazel/labgrid/strategy/transition.py @@ -1,13 +1,32 @@ from contextlib import contextmanager +from dataclasses import dataclass +from os import environ + +from labgrid.strategy import Strategy + + +@dataclass(frozen=True) +class State: + desired: str + initial: str | None + final: str | None + + @staticmethod + def from_env() -> "State": + return State( + desired=environ["LG_STATE"], + initial=environ.get("LG_INITIAL_STATE", None), + final=environ.get("BZL_LG_FINAL_STATE", None), + ) @contextmanager -def transition(strategy, initial, desired, final): +def transition(strategy: Strategy, state: State): try: - if initial is not None: - strategy.force(initial) - strategy.transition(desired) + if state.initial is not None: + strategy.force(state.initial) + strategy.transition(state.desired) yield finally: - if final is not None: - strategy.transition(final) + if state.final is not None: + strategy.transition(state.final) diff --git a/bazel/python/runfiles/BUILD.bazel b/bazel/python/runfiles/BUILD.bazel new file mode 100644 index 0000000000000000000000000000000000000000..35526b81c21144c8fe9de069aede20333f2c0dab --- /dev/null +++ b/bazel/python/runfiles/BUILD.bazel @@ -0,0 +1,11 @@ +load("@rules_python//python:defs.bzl", "py_library") + +py_library( + name = "runfiles", + srcs = [ + "__init__.py", + "runfiles.py", + ], + visibility = ["//:__subpackages__"], + deps = ["@rules_python//python/runfiles"], +) diff --git a/bazel/python/runfiles/__init__.py b/bazel/python/runfiles/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..724ec35947baf337e07cd391898c84d2ce614e09 --- /dev/null +++ b/bazel/python/runfiles/__init__.py @@ -0,0 +1,3 @@ +from .runfiles import runfile, RunfileNotFoundError + +__all__ = ["runfile", "RunfileNotFoundError"] diff --git a/bazel/python/runfiles/runfiles.py b/bazel/python/runfiles/runfiles.py new file mode 100644 index 0000000000000000000000000000000000000000..3ebbb9a6de625d2c4441198f53721cb2768b51d1 --- /dev/null +++ b/bazel/python/runfiles/runfiles.py @@ -0,0 +1,15 @@ +from pathlib import Path + +from python.runfiles import Runfiles + + +class RunfileNotFoundError(FileNotFoundError): + pass + + +def runfile(path: Path) -> Path: + runfiles = Runfiles.Create() + resolved = runfiles.Rlocation(path) + if resolved and Path(resolved).exists(): + return Path(resolved) + raise RunfileNotFoundError(path) diff --git a/labgrid/executor/BUILD.bazel b/labgrid/executor/BUILD.bazel index fdc6f791341c92e504d04efdc2fd0f251ac6128b..1effed64887fe57273b021f473a8df7cd873c248 100644 --- a/labgrid/executor/BUILD.bazel +++ b/labgrid/executor/BUILD.bazel @@ -18,8 +18,8 @@ py_binary( visibility = ["//:__subpackages__"], deps = [ "//bazel/labgrid/executor", + "//bazel/python/runfiles", "//labgrid/config:managers", - "@rules_python//python/runfiles", ], ) @@ -29,7 +29,7 @@ py_library( data = ["hello-world.txt"], deps = [ "//bazel/labgrid/executor", - "@rules_python//python/runfiles", + "//bazel/python/runfiles", ], ) diff --git a/labgrid/executor/executor.py b/labgrid/executor/executor.py index 9cd9c1b0107e54b9080c7f1ade89359c00dc76aa..467d90276ba303fa00f60e063b8300f11ce4d7f8 100644 --- a/labgrid/executor/executor.py +++ b/labgrid/executor/executor.py @@ -12,21 +12,8 @@ from shutil import which from subprocess import CalledProcessError, run from sys import argv -from python.runfiles import Runfiles - from bazel.labgrid.executor.manager import Data, Manager - - -class RunfileNotFoundError(FileNotFoundError): - pass - - -def runfile(path: Path) -> Path: - runfiles = Runfiles.Create() - resolved = runfiles.Rlocation(path) - if resolved and Path(resolved).exists(): - return Path(resolved) - raise RunfileNotFoundError(path) +from bazel.python.runfiles import runfile, RunfileNotFoundError def resolve(value: str) -> str: diff --git a/labgrid/executor/host.py b/labgrid/executor/host.py index d966d24f84357bfd97aa7a2e75d7fab8d21bf653..0b3b913bb0bdda03303ffa8949c80162f79d9428 100644 --- a/labgrid/executor/host.py +++ b/labgrid/executor/host.py @@ -3,15 +3,14 @@ from __future__ import annotations from contextlib import contextmanager from typing import Iterator -from python.runfiles import Runfiles from bazel.labgrid.executor.manager import Data +from bazel.python.runfiles import runfile @contextmanager def manager(input_data: Data) -> Iterator[Data]: - runfiles = Runfiles.Create() - path = runfiles.Rlocation("_main/labgrid/executor/hello-world.txt") + path = runfile("_main/labgrid/executor/hello-world.txt") with open(path) as stream: data = stream.read() env = { diff --git a/labgrid/run/BUILD.bazel b/labgrid/run/BUILD.bazel index f55b312e8bd86d5a162809093c45249620c0afb2..09e5125f7bd098183631dd3e9141d503e3224f2c 100644 --- a/labgrid/run/BUILD.bazel +++ b/labgrid/run/BUILD.bazel @@ -1,36 +1,14 @@ -load("@rules_labgrid//labgrid/config:transition.bzl", "labgrid_config_transition") load("@rules_python//python:defs.bzl", "py_binary") -TOOLS = [ - "mktemp", - "mkdir", - "rm", - "env", - "mv", -] - -[ - labgrid_config_transition( - name = tool, - srcs = ["@ape//ape:{}".format(tool)], - ) - for tool in TOOLS -] - py_binary( name = "run", - srcs = ["run.py"], - data = [ - ":{}".format(tool) - for tool in TOOLS + srcs = [ + "__main__.py", + "run.py", ], + imports = ["."], + main = "__main__.py", tags = ["manual"], visibility = ["//visibility:public"], - deps = [ - "//bazel/labgrid/strategy", - "//bazel/labgrid/util:target", - "//labgrid:pkg", - "//labgrid/config:deps", - "@rules_python//python/runfiles", - ], + deps = ["//bazel/labgrid/runner"], ) diff --git a/labgrid/run/__main__.py b/labgrid/run/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..621254e3e65b39490b4b9c3674930ee803bf1c5f --- /dev/null +++ b/labgrid/run/__main__.py @@ -0,0 +1,131 @@ +#! /usr/bin/env python3 +from __future__ import annotations + +from argparse import REMAINDER, ArgumentParser, ArgumentTypeError +from os import environ, getenv +from pathlib import Path +from string import Template +from subprocess import CalledProcessError +from sys import argv, stderr + +from run import run + +from bazel.labgrid.runner import FileTransfer +from bazel.python.runfiles import RunfileNotFoundError, runfile + + +def arguments(prsr: ArgumentParser) -> None: + prsr.add_argument( + "program", + metavar="PROG", + help="The program to run on the device.", + type=resolve, + ) + prsr.add_argument( + "arguments", metavar="ARG", nargs=REMAINDER, help="Command to run over SSH." + ) + prsr.add_argument( + "--env", + help=( + "The environment variables to be passed to the target for execution of the program. " + "Suports VAR=VAL syntax or uses the host environment variable if a single key is" + "specified." + ), + action="append", + required=False, + default=[], + metavar="KEY[=VALUE]", + type=env, + ) + prsr.add_argument( + "--get", + help="Transfer files from target at end of execution. Relative paths are resolved to the execution root on both the local and remote. `REMOTE` can have a trailing `?` to indicate it's optional. `LOCAL` performs environment variable substitution. `LOCAL` can be omitted, which will use the same path for remote and local.", + type=get, + metavar="REMOTE[?][:LOCAL]", + action="append", + default=[], + dest="downloads", + ) + prsr.add_argument( + "--put", + help="Transfer files to the target before execution. Relative paths are resolved to the execution root on both the local and remote. `LOCAL` performs environment variable substitution. `REMOTE` can be omitted, which will use the same path for remote and local.", + type=put, + metavar="LOCAL[:REMOTE]", + action="append", + default=[], + dest="uploads", + ) + + +def resolve(value: str) -> Path: + try: + return runfile(value) + except RunfileNotFoundError: + return Path(value) + + +def env(arg: str) -> tuple[str, str]: + key, sep, value = arg.partition("=") + if not sep: + value = getenv(key) + if value is None: + raise ArgumentTypeError(f"Env var {key} does not exist in host environment") + return (key, value) + + +def get(value: str) -> FileTransfer: + # FIXME: Extract `BazelArgumentParser` to handle quoting + remote, found, local = value.removeprefix("'").removesuffix("'").partition(":") + optional = remote.endswith("?") + if optional: + remote = remote.removesuffix("?") + if found: + local = Template(local).substitute(environ) + else: + local = remote + return FileTransfer(remote, local, optional) + + +def put(value: str) -> FileTransfer: + # FIXME: Extract `BazelArgumentParser` to handle quoting + local, found, remote = value.removeprefix("'").removesuffix("'").partition(":") + if found: + local = Template(local).substitute(environ) + else: + remote = local + return FileTransfer(remote, local) + + +def main(exe: Path, *args: str) -> int: + prsr = ArgumentParser( + prog=str(exe), description="Runs a command on a LabGrid device." + ) + + arguments(prsr) + parsed = prsr.parse_args(args) + + try: + return run( + parsed.program, + *parsed.arguments, + downloads=parsed.downloads, + uploads=parsed.uploads, + env=dict(parsed.env), + ) + except CalledProcessError as e: + print(f"fatal: subprocess failed: {e}", file=stderr) + if e.stdout is not None: + print(e.stdout, file=stderr) + if e.stderr is not None: + print(e.stderr, file=stderr) + return e.returncode + except KeyboardInterrupt: + return 130 + + +def entry(): + exit(main(Path(argv[0]), *argv[1:])) + + +if __name__ == "__main__": + entry() diff --git a/labgrid/run/run.py b/labgrid/run/run.py index 53cd1b054a9700e6136edfbc239b18eb5b74fad1..e7a04dbcb3a3974f1429e2208a1ed718b7041ff2 100644 --- a/labgrid/run/run.py +++ b/labgrid/run/run.py @@ -1,403 +1,28 @@ -#! /usr/bin/env python3 from __future__ import annotations -from typing import Any -import logging -from argparse import REMAINDER, Action, ArgumentParser, ArgumentTypeError, Namespace -from dataclasses import dataclass, replace -from os import environ, linesep, getenv -from pathlib import Path, PurePath +from pathlib import Path from shlex import join -from string import Template -from subprocess import CalledProcessError -from sys import argv, stderr, stdout from typing import ( - Final, Iterator, - MutableMapping, - Optional, - Sequence, - Text, - TypeVar, - Union, + Mapping, ) -from python.runfiles import Runfiles +from bazel.labgrid.runner import FileTransfer, runner -from bazel.labgrid.strategy import transition -from bazel.labgrid.util.target import TemporaryDirectory -from labgrid import Environment -from labgrid.logging import basicConfig, StepLogger -from labgrid.driver.exception import ExecutionError -T = TypeVar("T") -logger = logging.getLogger(__name__) -LOG_DEFAULT = logging.ERROR - - -class RunfileNotFoundError(FileNotFoundError): - pass - - -def runfile(path: str) -> Path: - runfiles = Runfiles.Create() - resolved = Path(runfiles.Rlocation(path)) - if not resolved.exists(): - raise RunfileNotFoundError(f"runfile not found: {path}") - return resolved - - -@dataclass -class Tools: - mktemp: Path = runfile("ape/ape/assimilate/mktemp/mktemp") - rm: Path = runfile("ape/ape/assimilate/rm/rm") - mkdir: Path = runfile("ape/ape/assimilate/mkdir/mkdir") - env: Path = runfile("ape/ape/assimilate/env/env") - mv: Path = runfile("ape/ape/assimilate/mv/mv") - - -def resolve(value: str) -> Path: - try: - return runfile(value) - except RunfileNotFoundError: - return Path(value) - - -@dataclass(frozen=True) -class State: - desired: str - initial: str | None = None - final: str | None = None - - -class StateAction(Action): - __default: Final[MutableMapping] = {} - - def __init__( - self, - option_strings: Sequence[str], - dest: str, - default: str | None, - const: Optional[T] = None, - help: Optional[str] = None, - ): - dest, found, attr = dest.partition(".") - if not found: - raise ValueError("`dest` must contain `.`") - self.attr = attr - super().__init__( - option_strings, - dest, - nargs=0, - const=const, - default=default, - type=None, - choices=None, - required=False, - help=help, - metavar=None, - ) - - @property - def default(self) -> State: - return State(**self.__default) - - @default.setter - def default(self, value: str | None) -> None: - self.__default[self.attr] = value - - def __call__( - self, - parser: ArgumentParser, - namespace: Namespace, - values: Union[str, Sequence[Any], None], - option_string: Optional[Text] = None, - ) -> None: - assert isinstance(values, str) - try: - state = getattr(namespace, self.dest) - except AttributeError: - state = self.default - state = replace(state, **{self.attr: values}) - setattr(namespace, self.dest, state) - - -class LogLevelAction(Action): - def __call__( - self, - parser: ArgumentParser, - namespace: Namespace, - values: Union[str, Sequence[Any], None], - option_string: Optional[Text] = None, - ) -> None: - assert isinstance(values, str) - setattr(namespace, self.dest, LogLevelAction.coerce(values)) - - def coerce(value: str) -> int: - # The Logging module converts strings to level integers with this (unintuitively named) function - log_level = logging.getLevelName(value) - if isinstance(log_level, str): - # getLevelName will return a string if the level was not recognised - raise ArgumentTypeError(f"Unsupported verbosity envvar: '{value}'") - - return log_level - - -def env(arg: str) -> str: - key, sep, value = arg.partition("=") - if not sep: - value = getenv(key) - - if value is None: - raise ArgumentTypeError(f"Env var {key} does not exist in host environment") - else: - return f"{key}={value}" - - return arg - - -def arguments(prsr: ArgumentParser) -> None: - prsr.add_argument( - "program", - metavar="PROG", - help="The program to run on the device.", - type=resolve, - ) - prsr.add_argument( - "arguments", metavar="ARG", nargs=REMAINDER, help="Command to run over SSH." - ) - prsr.add_argument( - "--config", - help="The LabGrid configuration.", - default=PurePath(environ["LG_ENV"]), - type=PurePath, - ) - prsr.add_argument( - "--initial-state", - help="The state to transition the LabGrid strategy to initially", - dest="state.initial", - action=StateAction, - default=environ.get("LG_INITIAL_STATE", None), - ) - prsr.add_argument( - "--state", - help="The state to transition the LabGrid strategy to", - dest="state.desired", - action=StateAction, - default=environ["LG_STATE"], - ) - prsr.add_argument( - "--final-state", - help="The state to transition the LabGrid strategy to at the end", - dest="state.final", - action=StateAction, - default=environ.get("BZL_LG_FINAL_STATE", None), - ) - prsr.add_argument( - "--env", - help=( - "The environment variables to be passed to the target for execution of the program. " - "Suports VAR=VAL syntax or uses the host environment variable if a single key is" - "specified." - ), - action="append", - required=False, - default=[], - metavar="KEY[=VALUE]", - type=env, - ) - prsr.add_argument( - "--get", - help="Transfer files from target at end of execution. Relative paths are resolved to the execution root on both the local and remote. `REMOTE` can have a trailing `?` to indicate it's optional. `LOCAL` performs environment variable substitution. `LOCAL` can be omitted, which will use the same path for remote and local.", - type=get, - metavar="REMOTE[?][:LOCAL]", - action="append", - default=[], - ) - prsr.add_argument( - "--put", - help="Transfer files to the target before execution. Relative paths are resolved to the execution root on both the local and remote. `LOCAL` performs environment variable substitution. `REMOTE` can be omitted, which will use the same path for remote and local.", - type=put, - metavar="LOCAL[:REMOTE]", - action="append", - default=[], - ) - prsr.add_argument( - "--log-level", - action=LogLevelAction, - help=f"Verbosity level of logging. Uses standard python logging levels (e.g. WARNING, INFO, DEBUG), the default is {logging.getLevelName(LOG_DEFAULT)}. Can also be set with LOG_LEVEL environment variable", - default=LogLevelAction.coerce( - environ.get("LOG_LEVEL", logging.getLevelName(LOG_DEFAULT)) - ), - dest="log_level", - ) - - -def get(value: str) -> FileTransfer: - # FIXME: Extract `BazelArgumentParser` to handle quoting - remote, found, local = value.removeprefix("'").removesuffix("'").partition(":") - optional = remote.endswith("?") - if optional: - remote = remote.removesuffix("?") - if found: - local = Template(local).substitute(environ) - else: - local = remote - return FileTransfer(PurePath(remote), PurePath(local), optional) - - -def put(value: str) -> FileTransfer: - # FIXME: Extract `BazelArgumentParser` to handle quoting - local, found, remote = value.removeprefix("'").removesuffix("'").partition(":") - if found: - local = Template(local).substitute(environ) - else: - remote = local - return FileTransfer(PurePath(remote), PurePath(local), False) - - -@dataclass(frozen=True) -class FileTransfer: - remote: PurePath - local: PurePath - optional: bool - - def resolve(self, remote: PurePath) -> FileTransfer: - return FileTransfer( - remote=self.join(remote, self.remote), - local=self.join(Path.cwd(), self.local), - optional=self.optional, - ) - - @staticmethod - def join(base: PurePath, path: PurePath) -> PurePath: - if not path.is_absolute(): - return base / path - return path - - -# TODO: move this out to `bazel.labgrid.run` on day to allow downstream to re-use def run( program: Path, *arguments: str, - config: Path, - state: State, - get: Iterator[FileTransfer], - put: Iterator[FileTransfer], - env: Iterator[str], - log_level: int, - tools: Tools = Tools(), + downloads: Iterator[FileTransfer], + uploads: Iterator[FileTransfer], + env: Mapping[str, str], ) -> int: - runfiles_dir = None - # Allow any Bazel binaries used in the configuration to work - try: - # FIXME: Convert into context manager - runfiles_dir = environ["RUNFILES_DIR"] - del environ["RUNFILES_DIR"] - except KeyError: - pass - - # Activating the StepLogger requires a DEBUG level of verbosity - basicConfig(level=log_level) - if log_level <= logging.DEBUG: - StepLogger.start() - - # Start up Docker container with LabGrid - target = Environment(str(config)).get_target() - strategy = target.get_driver("Strategy") - with transition(strategy, state.initial, state.desired, state.final): - # Retrieve the communication protocols - shell = target.get_driver("CommandProtocol") - transfer = target.get_driver("FileTransferProtocol") - with TemporaryDirectory( - shell, - transfer, - tools.mktemp, - tools.rm, - tools.mkdir, - tools.mv, - ) as temp: - # Transfer the provided program over to the device - remote = temp / program.name - transfer.put(f"{program}", f"{remote}") - - # Transfer program runfiles - # TODO: implement a `BazelFileTransfer` that transfers runfiles implicitly in `put` - # Make that the default file transfer implementation so _any_ file transfer automatically gets the runfiles as well - src_runfiles = program.with_suffix(".runfiles") - dest_runfiles = remote.with_suffix(".runfiles") - if src_runfiles.exists(): - transfer.put(f"{src_runfiles}", f"{dest_runfiles}") - elif runfiles_dir is not None: - transfer.put(f"{runfiles_dir}", f"{dest_runfiles}") - - # Transfer 'env' binary over to the device - transfer.put(f"{tools.env}", f"{temp / tools.env.name}") - - # Transfer user defined files - for unresolved in put: - resolved = unresolved.resolve(temp) - temp.mkdir(resolved.remote.parent) - transfer.put(resolved.local, resolved.remote) - - # Construct the CLI using env, the env vars and then the program and args. - cmd = join((f"{remote}", *arguments)) - cmd = f"cd {temp} && {temp / tools.env.name} {' '.join(env)} {cmd}" - - # Run the transferred program - out, err, code = shell.run(cmd) - for line in out: - stdout.write(f"{line}{linesep}") - for line in err: - stderr.write(f"{line}{linesep}") - - # Copy the files back from the target - - for unresolved in get: - resolved = unresolved.resolve(temp) - # FIXME: Check if file exists with `@ape//ape:test` instead of assuming any error is due to a missing file - try: - transfer.get(resolved.remote, resolved.local) - except ExecutionError as e: - if not resolved.optional and code == 0: - raise e - - return code - - -def main(exe: Path, *args: str) -> int: - prsr = ArgumentParser( - prog=str(exe), description="Runs a command on a LabGrid device." - ) - - arguments(prsr) - parsed = prsr.parse_args(args) - - try: - return run( - parsed.program, - *parsed.arguments, - state=parsed.state, - config=parsed.config, - get=parsed.get, - put=parsed.put, - env=parsed.env, - log_level=parsed.log_level, - ) - except CalledProcessError as e: - print(f"fatal: subprocess failed: {e}", file=stderr) - if e.stdout is not None: - print(e.stdout, file=stderr) - if e.stderr is not None: - print(e.stderr, file=stderr) - return e.returncode - except KeyboardInterrupt: - return 130 - - -def entry(): - exit(main(Path(argv[0]), *argv[1:])) + with runner() as r: + uploads += [FileTransfer(program.name, program, include_runfiles=True)] + cmd = join((f"./{program.name}", *arguments)) + r.put(uploads) + code = r.run(cmd, env) + r.get(downloads, code) -if __name__ == "__main__": - entry() + return code