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

from __future__ import annotations

import asyncio
import grp
import os
import pwd
import sys

from functools import cache
from typing import TYPE_CHECKING, Iterable, override

from ..base import Result, StatResult
from ..ExecContext import ExecContext as Base
from ..log import ERR, NOTICE, log

if TYPE_CHECKING:
    from typing import Any, Callable

    from ..Uri import Uri

class Local(Base):

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

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

    @override
    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: int, *args: Any, verbose: bool = verbose) -> None:
            if verbose:
                log(prio, log_prefix, *args)

        def __make_pty_reader(
            collector: list[bytes],
            enc_for_verbose: str,
        ) -> Callable[[int], bytes]:

            def _read(fd: int) -> bytes:
                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() -> int:
                    # 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, cmd = cmd)

            # -- 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: Any, prio: int, collector: list[bytes], log_enc: str
            ) -> None:
                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 (cmd_input is not None and stdin is asyncio.subprocess.PIPE
                    and proc.stdin is not None):
                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, cmd = cmd)

        finally:
            if cwd is not None:
                os.chdir(cwd)

    @override
    async def _unlink(self, path: str) -> None:
        os.unlink(path)

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

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

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

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

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

    @override
    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)

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

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

    @override
    async def _readdir(self, path: str) -> Iterable[str]:
        return os.listdir(path)