|
| 1 | +"""tee like run implementation.""" |
| 2 | +# This file is a modified version of https://github.com/pycontribs/subprocess-tee/blob/daffcbbf49fc5a2c7f3eaf75551f08fac0b9b63d/src/subprocess_tee/__init__.py |
| 3 | +# |
| 4 | +# It is licensed under the following license: |
| 5 | +# |
| 6 | +# The MIT License |
| 7 | +# Copyright (c) 2020 Sorin Sbarnea |
| 8 | +# Permission is hereby granted, free of charge, to any person obtaining a copy |
| 9 | +# of this software and associated documentation files (the "Software"), to deal |
| 10 | +# in the Software without restriction, including without limitation the rights |
| 11 | +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 12 | +# copies of the Software, and to permit persons to whom the Software is |
| 13 | +# furnished to do so, subject to the following conditions: |
| 14 | +# The above copyright notice and this permission notice shall be included in |
| 15 | +# all copies or substantial portions of the Software. |
| 16 | +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 17 | +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 18 | +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 19 | +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 20 | +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 21 | +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |
| 22 | +# THE SOFTWARE. |
| 23 | +import asyncio |
| 24 | +import os |
| 25 | +import platform |
| 26 | +import subprocess |
| 27 | +import sys |
| 28 | +from asyncio import StreamReader |
| 29 | +from typing import Any |
| 30 | +from typing import Callable |
| 31 | +from typing import Dict |
| 32 | +from typing import List |
| 33 | +from typing import Optional |
| 34 | +from typing import TYPE_CHECKING |
| 35 | +from typing import Union |
| 36 | + |
| 37 | +if TYPE_CHECKING: |
| 38 | + CompletedProcess = subprocess.CompletedProcess[Any] # pylint: disable=E1136 |
| 39 | +else: |
| 40 | + CompletedProcess = subprocess.CompletedProcess |
| 41 | + |
| 42 | +try: |
| 43 | + from shlex import join # type: ignore |
| 44 | +except ImportError: |
| 45 | + from subprocess import list2cmdline as join # pylint: disable=ungrouped-imports |
| 46 | + |
| 47 | + |
| 48 | +STREAM_LIMIT = 2 ** 23 # 8MB instead of default 64kb, override it if you need |
| 49 | + |
| 50 | + |
| 51 | +async def _read_stream(stream: StreamReader, callback: Callable[..., Any]) -> None: |
| 52 | + while True: |
| 53 | + line = await stream.readline() |
| 54 | + if line: |
| 55 | + callback(line) |
| 56 | + else: |
| 57 | + break |
| 58 | + |
| 59 | + |
| 60 | +async def _stream_subprocess(args: str, **kwargs: Any) -> CompletedProcess: |
| 61 | + platform_settings: Dict[str, Any] = {} |
| 62 | + if platform.system() == "Windows": |
| 63 | + platform_settings["env"] = os.environ |
| 64 | + |
| 65 | + # this part keeps behavior backwards compatible with subprocess.run |
| 66 | + tee = kwargs.get("tee", True) |
| 67 | + stdout = kwargs.get("stdout", sys.stdout) |
| 68 | + if stdout == subprocess.DEVNULL or not tee: |
| 69 | + stdout = open(os.devnull, "w") |
| 70 | + stderr = kwargs.get("stderr", sys.stderr) |
| 71 | + if stderr == subprocess.DEVNULL or not tee: |
| 72 | + stderr = open(os.devnull, "w") |
| 73 | + |
| 74 | + # We need to tell subprocess which shell to use when running shell-like |
| 75 | + # commands. |
| 76 | + # * SHELL is not always defined |
| 77 | + # * /bin/bash does not exit on alpine, /bin/sh seems bit more portable |
| 78 | + if "executable" not in kwargs and isinstance(args, str) and " " in args: |
| 79 | + platform_settings["executable"] = os.environ.get("SHELL", "/bin/sh") |
| 80 | + |
| 81 | + # pass kwargs we know to be supported |
| 82 | + for arg in ["cwd", "env"]: |
| 83 | + if arg in kwargs: |
| 84 | + platform_settings[arg] = kwargs[arg] |
| 85 | + |
| 86 | + # Some users are reporting that default (undocumented) limit 64k is too |
| 87 | + # low |
| 88 | + process = await asyncio.create_subprocess_shell( |
| 89 | + args, |
| 90 | + limit=STREAM_LIMIT, |
| 91 | + stdin=kwargs.get("stdin", False), |
| 92 | + stdout=asyncio.subprocess.PIPE, |
| 93 | + stderr=asyncio.subprocess.PIPE, |
| 94 | + **platform_settings, |
| 95 | + ) |
| 96 | + out: List[str] = [] |
| 97 | + err: List[str] = [] |
| 98 | + |
| 99 | + def tee_func(line: bytes, sink: List[str], pipe: Optional[Any]) -> None: |
| 100 | + line_str = line.decode("utf-8").rstrip() |
| 101 | + sink.append(line_str) |
| 102 | + if not kwargs.get("quiet", False): |
| 103 | + # This is modified from the default implementation since |
| 104 | + # we want all output to be interleved on the same stream |
| 105 | + print(line_str, file=sys.stderr) |
| 106 | + |
| 107 | + loop = asyncio.get_event_loop() |
| 108 | + tasks = [] |
| 109 | + if process.stdout: |
| 110 | + tasks.append( |
| 111 | + loop.create_task( |
| 112 | + _read_stream(process.stdout, lambda l: tee_func(l, out, stdout)) |
| 113 | + ) |
| 114 | + ) |
| 115 | + if process.stderr: |
| 116 | + tasks.append( |
| 117 | + loop.create_task( |
| 118 | + _read_stream(process.stderr, lambda l: tee_func(l, err, stderr)) |
| 119 | + ) |
| 120 | + ) |
| 121 | + |
| 122 | + await asyncio.wait(set(tasks)) |
| 123 | + |
| 124 | + # We need to be sure we keep the stdout/stderr output identical with |
| 125 | + # the ones procued by subprocess.run(), at least when in text mode. |
| 126 | + check = kwargs.get("check", False) |
| 127 | + stdout = None if check else "" |
| 128 | + stderr = None if check else "" |
| 129 | + if out: |
| 130 | + stdout = os.linesep.join(out) + os.linesep |
| 131 | + if err: |
| 132 | + stderr = os.linesep.join(err) + os.linesep |
| 133 | + |
| 134 | + return CompletedProcess( |
| 135 | + args=args, |
| 136 | + returncode=await process.wait(), |
| 137 | + stdout=stdout, |
| 138 | + stderr=stderr, |
| 139 | + ) |
| 140 | + |
| 141 | + |
| 142 | +def run(args: Union[str, List[str]], **kwargs: Any) -> CompletedProcess: |
| 143 | + """Drop-in replacement for subprocerss.run that behaves like tee. |
| 144 | + Extra arguments added by our version: |
| 145 | + echo: False - Prints command before executing it. |
| 146 | + quiet: False - Avoid printing output |
| 147 | + """ |
| 148 | + if isinstance(args, str): |
| 149 | + cmd = args |
| 150 | + else: |
| 151 | + # run was called with a list instead of a single item but asyncio |
| 152 | + # create_subprocess_shell requires command as a single string, so |
| 153 | + # we need to convert it to string |
| 154 | + cmd = join(args) |
| 155 | + |
| 156 | + check = kwargs.get("check", False) |
| 157 | + |
| 158 | + if kwargs.get("echo", False): |
| 159 | + # This is modified from the default implementation since |
| 160 | + # we want all output to be interleved on the same stream |
| 161 | + print(f"COMMAND: {cmd}", file=sys.stderr) |
| 162 | + |
| 163 | + loop = asyncio.get_event_loop() |
| 164 | + result = loop.run_until_complete(_stream_subprocess(cmd, **kwargs)) |
| 165 | + |
| 166 | + if check and result.returncode != 0: |
| 167 | + raise subprocess.CalledProcessError( |
| 168 | + result.returncode, cmd, output=result.stdout, stderr=result.stderr |
| 169 | + ) |
| 170 | + return result |
0 commit comments