structio/structio/parallel.py

"""
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