mirror of
ssh://git.janware.com/janware/proj/jw-pkg
synced 2026-04-24 17:23:36 +02:00
Jan Lindemann
888c0495ec
Add lib.base to provide basic definitions. For now, move the definiions of Result, Input and InputMode from ExecContext into lib.base. Having to import them from the ExecContect module is too heavy-handed for those simple types. Signed-off-by: Jan Lindemann <jan@janware.com>
148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import os, sys, subprocess, asyncio
|
|
|
|
from ..ExecContext import ExecContext as Base
|
|
from ..base import Result
|
|
|
|
from ..log import *
|
|
from ..util import pretty_cmd
|
|
|
|
class Local(Base):
|
|
|
|
def __init__(self, uri='local', *args, **kwargs) -> None:
|
|
super().__init__(uri, *args, **kwargs)
|
|
|
|
async def _run(
|
|
self,
|
|
cmd: list[str],
|
|
wd: str|None,
|
|
verbose: bool,
|
|
cmd_input: bytes|None,
|
|
env: dict[str, str]|None,
|
|
interactive: bool,
|
|
log_prefix: str
|
|
) -> Result:
|
|
|
|
def __log(prio, *args, verbose=verbose):
|
|
if verbose:
|
|
log(prio, log_prefix, *args)
|
|
|
|
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
|
|
|
|
cwd: str|None = None
|
|
if wd is not None:
|
|
cwd = os.getcwd()
|
|
os.chdir(wd)
|
|
|
|
try:
|
|
|
|
# -- interactive mode
|
|
|
|
if interactive:
|
|
|
|
import pty
|
|
|
|
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(cmd, master_read=reader)
|
|
finally:
|
|
os.environ.clear()
|
|
os.environ.update(old_env)
|
|
return pty.spawn(cmd, master_read=reader)
|
|
|
|
stdout_chunks: list[bytes] = []
|
|
enc_for_verbose = sys.stdout.encoding or "utf-8"
|
|
reader = __make_pty_reader(stdout_chunks, enc_for_verbose)
|
|
|
|
exit_code = await asyncio.to_thread(_spawn)
|
|
|
|
# PTY merges stdout/stderr
|
|
stdout = b"".join(stdout_chunks) if stdout_chunks else None
|
|
return Result(stdout, 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(
|
|
*cmd,
|
|
stdin=stdin,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
env=env,
|
|
)
|
|
|
|
stdout_parts: list[bytes] = []
|
|
stderr_parts: list[bytes] = []
|
|
|
|
# -- decoding for verbose output in pipe mode
|
|
stdout_log_enc = sys.stdout.encoding or "utf-8"
|
|
stderr_log_enc = sys.stderr.encoding or "utf-8"
|
|
|
|
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, stdout_log_enc)
|
|
),
|
|
asyncio.create_task(
|
|
read_stream(proc.stderr, ERR, stderr_parts, stderr_log_enc)
|
|
),
|
|
]
|
|
|
|
if stdin is asyncio.subprocess.PIPE:
|
|
proc.stdin.write(cmd_input)
|
|
await proc.stdin.drain()
|
|
proc.stdin.close()
|
|
|
|
exit_code = await proc.wait()
|
|
await asyncio.gather(*tasks)
|
|
|
|
stdout = b"".join(stdout_parts) if stdout_parts else None
|
|
stderr = b"".join(stderr_parts) if stderr_parts else None
|
|
|
|
return Result(stdout, stderr, exit_code)
|
|
|
|
finally:
|
|
if cwd is not None:
|
|
os.chdir(cwd)
|
|
|
|
async def _sudo(self, cmd: list[str], mod_env: dict[str, str], opts: list[str], *args, **kwargs) -> Result:
|
|
env: dict[str, str]|None = None
|
|
cmd_input: bytes|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)
|
|
return await self._run(cmdline, *args, **kwargs)
|