""" Module inspired by subprocess which allows for asynchronous multiprocessing """ import os import structio import platform import subprocess from subprocess import CalledProcessError, CompletedProcess, DEVNULL, PIPE from structio.io import FileStream from structio.core.syscalls import checkpoint if platform.system() == "Windows": # Windows doesn't really support non-blocking file # descriptors (except sockets), so we just use threads from structio.io.files import AsyncFile as FileStream class Process: """ Class similar to subprocess.Popen, but async. The constructor is analogous to its synchronous counterpart """ def __init__(self, *args, **kwargs): if "universal_newlines" in kwargs: # Not sure why? But everyone else is doing it so :shrug: raise RuntimeError("universal_newlines is not supported") if stdin := kwargs.get("stdin"): if stdin not in {PIPE, DEVNULL}: # Curio mentions stuff breaking if the child process # is passed a stdin fd that is set to non-blocking mode if hasattr(os, "set_blocking"): os.set_blocking(stdin.fileno(), True) # Delegate to Popen's constructor self._process: subprocess.Popen | None = None self._args = args self._kwargs = kwargs self.stdin = None self.stdout = None self.stderr = None self.returncode = None self.pid = -1 self._taskid = None async def terminate(self): """ Terminates the subprocess asynchronously """ return await structio.thread.run_in_worker( self._process.terminate, cancellable=True ) def start(self): """ Begin execution of the process """ self._process = subprocess.Popen(*self._args, **self._kwargs) self.pid = self._process.pid if self._process.stdin: self.stdin = FileStream(self._process.stdin) if self._process.stdout: self.stdout = FileStream(self._process.stdout) if self._process.stderr: self.stderr = FileStream(self._process.stderr) # self._taskid = structio.current_loop().add_shutdown_task(self.wait) async def is_running(self): """ Returns whether the process is currently running """ if self._process is None: return False elif self._process.poll() is None: return False return True async def wait(self): """ Async equivalent of subprocess.Popen.wait() """ if self._process is None: raise RuntimeError("process is not running yet") status = self._process.poll() if status is None: status = await structio.thread.run_in_worker( self._process.wait, cancellable=True ) self.returncode = status if self._taskid is not None: structio.current_loop().remove_shutdown_task(self._taskid) return status async def communicate(self, input=b"") -> tuple[bytes, bytes]: async with structio.create_pool() as pool: stdout = pool.spawn(self.stdout.readall) if self.stdout else None stderr = pool.spawn(self.stderr.readall) if self.stderr else None if input: await self.stdin.write(input) await self.stdin.close() # Awaiting a task object waits for its completion and # returns its return value! out = b"" err = b"" if stdout: out = await stdout if stderr: err = await stderr return out, err async def __aenter__(self): self.start() return self async def __aexit__(self, *args): if self.stdin: await self.stdin.close() if self.stdout: await self.stdout.close() if self.stderr: await self.stderr.close() await self.wait() def __getattr__(self, item): # Delegate to internal process object return getattr(self._process, item) async def run( args, *, stdin=None, input=None, stdout=None, stderr=None, shell=False, check=False ): """ Async version of subprocess.run() """ if input: stdin = subprocess.PIPE async with Process( args, stdin=stdin, stdout=stdout, stderr=stderr, shell=shell ) as process: try: stdout, stderr = await process.communicate(input) except: process.kill() raise if check and process.returncode: raise CalledProcessError(process.returncode, process.args, output=stdout, stderr=stderr) return CompletedProcess(process.args, process.returncode, stdout, stderr) async def check_output(args, *, stdin=None, stderr=None, shell=False, input=None): """ Async version of subprocess.check_output """ out = await run( args, stdout=PIPE, stdin=stdin, stderr=stderr, shell=shell, check=True, input=input, ) return out.stdout