structio/structio/parallel.py

174 lines
5.1 KiB
Python
Raw Normal View History

"""
Module inspired by subprocess which allows for asynchronous
multiprocessing
"""
import os
import structio
2023-08-23 21:56:25 +02:00
import platform
import subprocess
2023-06-19 17:34:44 +02:00
from subprocess import CalledProcessError, CompletedProcess, DEVNULL, PIPE
from structio.io import FileStream
from structio.core.syscalls import checkpoint
2024-02-23 13:11:14 +01:00
2023-08-23 21:56:25 +02:00
if platform.system() == "Windows":
# Windows doesn't really support non-blocking file
# descriptors (except sockets), so we just use threads
2023-08-23 21:56:25 +02:00
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
2023-08-23 21:56:25 +02:00
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:
2023-06-19 17:34:44 +02:00
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)
2023-06-19 17:34:44 +02:00
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(
2023-06-19 17:34:44 +02:00
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
"""
2023-06-19 17:34:44 +02:00
out = await run(
args,
stdout=PIPE,
stdin=stdin,
stderr=stderr,
shell=shell,
check=True,
input=input,
)
return out.stdout