pymake

A build system based on Build Systems à la Carte
git clone https://git.grace.moe/pymake
Log | Files | Refs | README

exec.py (2318B)


      1 from asyncio import create_subprocess_exec, create_subprocess_shell, gather
      2 from asyncio.streams import StreamReader, StreamWriter
      3 from asyncio.subprocess import Process, PIPE
      4 from collections import namedtuple
      5 import sys
      6 
      7 
      8 class ShellResult(namedtuple("ShellResult", "stdout stderr returncode")):
      9     __slots__ = ()
     10 
     11     @property
     12     def utf8stdout(self):
     13         return self.stdout.decode("utf-8")
     14 
     15     @property
     16     def utf8stderr(self):
     17         return self.stderr.decode("utf-8")
     18 
     19 
     20 EchoNothing = 0
     21 EchoStdout = 1
     22 EchoStderr = 2
     23 EchoAll = 3
     24 
     25 
     26 async def _exec_writer(
     27     proc: Process,
     28     ostream: StreamWriter,
     29     input: bytes | bytearray | memoryview | None = None,
     30 ):
     31     if input is not None:
     32         ostream.write(input)
     33         await ostream.drain()
     34     return await proc.wait()
     35 
     36 
     37 async def _exec_reader(istream: StreamReader, ostream=None):
     38     contents = b""
     39     while not istream.at_eof():
     40         chunk = await istream.read(4096 * 16)
     41         contents += chunk
     42         if ostream:
     43             ostream.write(chunk)
     44             ostream.flush()
     45     return contents
     46 
     47 
     48 async def communicate_echo_wait(
     49     proc: Process,
     50     input: bytes | bytearray | memoryview | None = None,
     51     echo: int = EchoNothing,
     52 ) -> ShellResult:
     53     stdout, stderr, returncode = await gather(
     54         _exec_reader(
     55             proc.stdout,  # type: ignore
     56             sys.stdout.buffer if echo & EchoStdout else None,
     57         ),
     58         _exec_reader(
     59             proc.stderr,  # type: ignore
     60             sys.stderr.buffer if echo & EchoStderr else None,
     61         ),
     62         _exec_writer(
     63             proc,
     64             proc.stdin,  # type: ignore
     65             input,
     66         ),
     67     )
     68     return ShellResult(stdout, stderr, returncode)
     69 
     70 
     71 async def exec(
     72     prog,
     73     *args,
     74     input: bytes | bytearray | memoryview | None = None,
     75     echo: int = EchoNothing,
     76 ) -> ShellResult:
     77     return await communicate_echo_wait(
     78         await create_subprocess_exec(prog, *args, stdin=PIPE, stdout=PIPE, stderr=PIPE),
     79         input,
     80         echo,
     81     )
     82 
     83 
     84 async def shell(
     85     cmd,
     86     input: bytes | bytearray | memoryview | None = None,
     87     echo: int = EchoNothing,
     88 ) -> ShellResult:
     89     return await communicate_echo_wait(
     90         await create_subprocess_shell(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE),
     91         input,
     92         echo,
     93     )