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 )