# -*- coding: utf-8 -*- import os, sys, subprocess, asyncio from ..util import run_cmd from ..ExecContext import ExecContext as Base from ..ExecContext import Result from ..log import * from ..util import pretty_cmd class Local(Base): async def _run( self, args: list[str], wd: str|None = None, throw: bool = True, verbose: bool = False, cmd_input: str|None = None, env: dict[str, str]|None = None, title: str=None, output_encoding: str|None = None, # None => unchanged; "bytes" => return raw bytes ) -> tuple[str|bytes|None, str|bytes|None]: want_bytes = (output_encoding == "bytes") def __log(prio, *args, verbose=verbose): if verbose: log(prio, "|", *args) def __check_exit_code(code: int, stdout=None, stderr=None): if code == 0: return if (throw or verbose): msg = f'Command returned error {code}: {pretty_cmd(args, wd)}' if stderr: msg += ': ' + stderr.strip() if throw: raise RuntimeError(msg) def __make_pty_reader(collector: list[bytes], enc_for_verbose: str): def _read(fd): ret = os.read(fd, 1024) if not ret: return ret collector.append(ret) return ret return _read interactive = ( cmd_input == "mode:interactive" or (cmd_input == "mode:auto" and sys.stdin.isatty()) ) if verbose: delim_len = 120 delim = title if title is not None else f'---- Running {pretty_cmd(args, wd)} -' if interactive: log(NOTICE, delim) else: delim += '-' * max(0, delim_len - len(delim)) log(NOTICE, ',' + delim + ' >') cwd: str|None = None if wd is not None: cwd = os.getcwd() os.chdir(wd) try: # -- interactive mode if interactive: import pty stdout_chunks_b: list[bytes] = [] enc_for_verbose = ( (sys.stdout.encoding or "utf-8") if output_encoding in (None, "bytes") else output_encoding ) reader = __make_pty_reader(stdout_chunks_b, enc_for_verbose) def _spawn(): # Apply env in PTY mode by temporarily updating os.environ around spawn. if env: old_env = os.environ.copy() try: os.environ.update(env) return pty.spawn(args, master_read=reader) finally: os.environ.clear() os.environ.update(old_env) return pty.spawn(args, master_read=reader) exit_code = await asyncio.to_thread(_spawn) __check_exit_code(exit_code) # PTY merges stdout/stderr stdout_b = b"".join(stdout_chunks_b) if stdout_chunks_b else None if want_bytes: return stdout_b, None, exit_code stdout_dec_enc = (sys.stdout.encoding or "utf-8") if output_encoding is None else output_encoding stdout_s = stdout_b.decode(stdout_dec_enc, errors="replace") if stdout_b is not None else None return stdout_s, None, exit_code # -- non-interactive mode stdin = ( asyncio.subprocess.DEVNULL if cmd_input is None else asyncio.subprocess.PIPE ) proc = await asyncio.create_subprocess_exec( *args, stdin=stdin, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env, ) stdout_parts_b: list[bytes] = [] stderr_parts_b: list[bytes] = [] # -- decoding for verbose output in pipe mode if output_encoding is None or want_bytes: stdout_log_enc = sys.stdout.encoding or "utf-8" stderr_log_enc = sys.stderr.encoding or "utf-8" else: stdout_log_enc = output_encoding stderr_log_enc = output_encoding async def read_stream(stream, prio, collector: list[bytes], log_enc: str): buf = b"" while True: chunk = await stream.read(4096) if not chunk: break collector.append(chunk) if verbose: buf += chunk while b"\n" in buf: line, buf = buf.split(b"\n", 1) __log(prio, line.decode(log_enc, errors="replace")) if verbose and buf: # flush trailing partial line (no newline) __log(prio, buf.decode(log_enc, errors="replace")) tasks = [ asyncio.create_task( read_stream(proc.stdout, NOTICE, stdout_parts_b, stdout_log_enc) ), asyncio.create_task( read_stream(proc.stderr, ERR, stderr_parts_b, stderr_log_enc) ), ] if stdin is asyncio.subprocess.PIPE: proc.stdin.write(cmd_input.encode(sys.stdout.encoding or "utf-8")) await proc.stdin.drain() proc.stdin.close() exit_code = await proc.wait() await asyncio.gather(*tasks) if want_bytes: __check_exit_code(exit_code) stdout_b = b"".join(stdout_parts_b) if stdout_parts_b else None stderr_b = b"".join(stderr_parts_b) if stderr_parts_b else None if want_bytes: return stdout_b, stderr_b, exit_code if output_encoding is None: stdout_dec_enc = sys.stdout.encoding or "utf-8" stderr_dec_enc = sys.stderr.encoding or "utf-8" else: stdout_dec_enc = output_encoding stderr_dec_enc = output_encoding stdout_s = stdout_b.decode(stdout_dec_enc, errors="replace") if stdout_b is not None else None stderr_s = stderr_b.decode(stderr_dec_enc, errors="replace") if stderr_b is not None else None if not want_bytes: __check_exit_code(exit_code, stdout=stdout_s, stderr=stderr_s) return stdout_s, stderr_s, exit_code finally: if cwd is not None: os.chdir(cwd) if verbose and not interactive: log(NOTICE, '`' + delim + ' <') async def _sudo(self, cmd: list[str], mod_env: dict[str, str] = {}, opts: list[str]=[], verbose=True) -> Result: env: dict[str, str]|None = None cmd_input: str|None = None if mod_env: env = os.environ.copy() env.update(mod_env) cmdline = [] if os.getuid() != 0: cmdline.append('/usr/bin/sudo') if env is not None: cmdline.append('--preserve-env=' + ','.join(mod_env.keys())) cmdline.extend(opts) cmdline.extend(cmd) if self.interactive: cmd_input = "mode:interactive" return await self._run(cmdline, throw=True, verbose=verbose, env=env, cmd_input=cmd_input)