167 lines
4.7 KiB
Python
167 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import io
|
|
import os
|
|
import runpy
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Callable, Iterable, Optional, Tuple
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def pushd(path: Path):
|
|
current = Path.cwd()
|
|
os.chdir(path)
|
|
try:
|
|
yield
|
|
finally:
|
|
os.chdir(current)
|
|
|
|
|
|
class StreamCapture:
|
|
def __init__(self, on_output: Optional[Callable[[str], None]], collect: io.StringIO) -> None:
|
|
self.on_output = on_output
|
|
self.collect = collect
|
|
self.buffer = ""
|
|
|
|
def write(self, text: str) -> int:
|
|
if text is None:
|
|
return 0
|
|
text = str(text)
|
|
self.collect.write(text)
|
|
if self.on_output:
|
|
text = text.replace("\r", "\n")
|
|
self.buffer += text
|
|
while "\n" in self.buffer:
|
|
line, self.buffer = self.buffer.split("\n", 1)
|
|
if line:
|
|
self.on_output(line)
|
|
return len(text)
|
|
|
|
def flush(self) -> None:
|
|
if self.on_output and self.buffer:
|
|
self.on_output(self.buffer)
|
|
self.buffer = ""
|
|
|
|
@property
|
|
def encoding(self) -> str:
|
|
return "utf-8"
|
|
|
|
@property
|
|
def errors(self) -> str:
|
|
return "replace"
|
|
|
|
def isatty(self) -> bool:
|
|
return False
|
|
|
|
def writable(self) -> bool:
|
|
return True
|
|
|
|
def fileno(self) -> int:
|
|
raise OSError("StreamCapture has no file descriptor")
|
|
|
|
|
|
def run_tool_module(
|
|
tool: str,
|
|
args: Iterable[str],
|
|
cwd: Path,
|
|
stream: bool,
|
|
on_output: Optional[Callable[[str], None]],
|
|
) -> Tuple[int, str]:
|
|
argv_backup = sys.argv[:]
|
|
prev_rich_disable = os.environ.get("RICH_DISABLE")
|
|
prev_rich_click_disable = os.environ.get("RICH_CLICK_DISABLE")
|
|
output = io.StringIO()
|
|
capture = StreamCapture(on_output if stream else None, output)
|
|
try:
|
|
sys.argv = [tool, *list(args)]
|
|
os.environ["RICH_DISABLE"] = "1"
|
|
os.environ["RICH_CLICK_DISABLE"] = "1"
|
|
with pushd(cwd), contextlib.redirect_stdout(capture), contextlib.redirect_stderr(capture):
|
|
try:
|
|
module_name = f"{tool}.__main__"
|
|
try:
|
|
runpy.run_module(module_name, run_name="__main__")
|
|
except ModuleNotFoundError as exc:
|
|
if exc.name in (module_name, tool):
|
|
runpy.run_module(tool, run_name="__main__")
|
|
else:
|
|
raise
|
|
code = 0
|
|
except SystemExit as exc:
|
|
code = int(exc.code) if isinstance(exc.code, int) else 0
|
|
except Exception:
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
code = 1
|
|
finally:
|
|
capture.flush()
|
|
sys.argv = argv_backup
|
|
if prev_rich_disable is None:
|
|
os.environ.pop("RICH_DISABLE", None)
|
|
else:
|
|
os.environ["RICH_DISABLE"] = prev_rich_disable
|
|
if prev_rich_click_disable is None:
|
|
os.environ.pop("RICH_CLICK_DISABLE", None)
|
|
else:
|
|
os.environ["RICH_CLICK_DISABLE"] = prev_rich_click_disable
|
|
return code, output.getvalue()
|
|
|
|
|
|
def run_esptool_api(
|
|
func: Callable[[], None],
|
|
stream: bool,
|
|
on_output: Optional[Callable[[str], None]],
|
|
) -> Tuple[int, str]:
|
|
output = io.StringIO()
|
|
capture = StreamCapture(on_output if stream else None, output)
|
|
try:
|
|
with contextlib.redirect_stdout(capture), contextlib.redirect_stderr(capture):
|
|
func()
|
|
code = 0
|
|
except SystemExit as exc:
|
|
code = int(exc.code) if isinstance(exc.code, int) else 1
|
|
except Exception:
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
code = 1
|
|
finally:
|
|
capture.flush()
|
|
return code, output.getvalue()
|
|
|
|
|
|
def run_command(
|
|
cmd: Iterable[str],
|
|
cwd: Path,
|
|
stream: bool = False,
|
|
on_output: Optional[Callable[[str], None]] = None,
|
|
) -> Tuple[int, str]:
|
|
if stream:
|
|
proc = subprocess.Popen(
|
|
list(cmd),
|
|
cwd=cwd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
output_parts = []
|
|
assert proc.stdout is not None
|
|
for line in proc.stdout:
|
|
output_parts.append(line)
|
|
if on_output:
|
|
on_output(line.rstrip("\n"))
|
|
proc.wait()
|
|
return proc.returncode, "".join(output_parts)
|
|
|
|
result = subprocess.run(list(cmd), cwd=cwd, capture_output=True, text=True)
|
|
combined = (result.stdout or "") + (result.stderr or "")
|
|
if on_output and combined:
|
|
for line in combined.splitlines():
|
|
on_output(line)
|
|
return result.returncode, combined
|