jw-pkg/src/python/jw/pkg/lib/ec/Local.py

# -*- coding: utf-8 -*-

from __future__ import annotations
from typing import TYPE_CHECKING

import os, sys, subprocess, asyncio, pwd, grp, stat

from functools import cache

from ..ExecContext import ExecContext as Base
from ..base import Result, StatResult

from ..log import *
from ..util import pretty_cmd

if TYPE_CHECKING:
    from ..Uri import Uri

class Local(Base):

    def __init__(self, uri: str|Uri='local', *args, **kwargs) -> None:
        super().__init__(uri, *args, **kwargs)

    @cache
    def _username(self) -> str:
        return pwd.getpwuid(os.getuid()).pw_name,

    async def _run(
        self,
        cmd: list[str],
        wd: str|None,
        verbose: bool,
        cmd_input: bytes|None,
        mod_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 mod_env:
                        old_env = os.environ.copy()
                        try:
                            os.environ.update(mod_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

            if mod_env:
                new_env = os.environ.copy()
                new_env.update(mod_env)
                mod_env = new_env

            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=stdin,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=mod_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 _unlink(self, path: str) -> None:
        os.unlink(path)

    async def _erase(self, path: str) -> None:
        if os.isdir(path):
            shutil.rmtree(path)
            return
        os.unlink(path)

    async def _rename(self, src: str, dst: str) -> None:
        os.rename(src, dst)

    async def _mkdir(self, name: str, mode: int) -> None:
        os.mkdir(name, mode)

    async def _stat(self, path: str, follow_symlinks: bool) -> StatResult:
        return StatResult.from_os(os.stat(path, follow_symlinks=follow_symlinks))

    async def _file_exists(self, path: str) -> bool:
        return os.path.exists(path)

    async def _chown(self, path: str, owner: str|None, group: str|None) -> None:
        uid = pwd.getpwnam(owner).pw_uid if owner else -1
        gid = grp.getgrnam(group).gr_gid if group else -1
        os.chown(path, uid, gid)

    async def _chmod(self, path: str, mode: int) -> None:
        os.chmod(path, mode)

    async def _is_dir(self, path: str, follow_symlinks: bool) -> bool:
        if (not follow_symlinks) and os.islink(path):
            return False
        return os.path.isdir(path)