blender-package-manager-addon/subprocess_adapter.py

import logging
from multiprocessing import Process, Pipe
from bpy.types import Operator

log = logging.getLogger(__name__)

def subprocess_operator(cls: Operator, polling_interval=.01) -> Operator:
    """
    Class decorator which wraps Operator methods with setup code for running a subprocess.

    Expects args for Process() to defined in cls.proc_args and cls.proc_kwargs. For example,
    setting cls.proc_kwargs = {'target:' some_func} can be used to run 'some_func' in
    a subprocess.
    """

    def decoratify(cls, methodname, decorator):
        """
        Decorate `cls.methodname` with `decorator` if method `methodname` exists in cls
        else just call the decorator with an empty function
        """
        orig_method = getattr(cls, methodname, None)
        if orig_method:
            setattr(cls, methodname, decorator(orig_method))
        else:
            # HACK: need a no-op which accepts any arguments
            setattr(cls, methodname, decorator(lambda *args, **kwargs: None))


    def decorate_execute(orig_execute):
        def execute(self, context):
            orig_execute(self, context)
            call_copy_of_method_if_exist(cls, 'execute', self, context)
            return self.invoke(context, None)
        return execute


    def decorate_invoke(orig_invoke):
        """
        Create pipe and modal timer, start subprocess
        """
        def invoke(self, context, event):
            self.pipe = Pipe() #HACK: do this first so operator-defined invoke can access it
            # this is needed because the operator is currently responsible for setting up the Process,
            # and the pipe is needed for the subprocess_function decorator.
            # TODO: Perhaps this responsibility should be handled here instead (simplifies things but looses some flexibility)

            orig_ret = orig_invoke(self, context, event)

            # HACK to allow early exits.
            # Perhaps subprocess_operators should define separately named methods, to avoid confusing mixing semantics in return types
            if not orig_ret.issubset({'RUNNING_MODAL'}):
                log.debug('Exiting early')
                return orig_ret

            self.proc = Process(
                    *getattr(self, 'proc_args', []),
                    **getattr(self, 'proc_kwargs', [])
                    )
            self.proc.start()
            # we have to explicitly close the end of the pipe we are NOT using,
            # otherwise no exception will be generated when the other process closes its end.
            self.pipe[1].close()

            wm = context.window_manager
            wm.modal_handler_add(self)
            self.timer = wm.event_timer_add(polling_interval, context.window)

            return {'RUNNING_MODAL'}
        return invoke

    def poll_subprocess(self):
        """
        Check the pipe for new messages (nonblocking). If there is a new message, call the callback
        """
        log.debug("polling")
        try:
            if self.pipe[0].poll():
                resp = self.pipe[0].recv()
                # assert(isinstance(resp, SubprocessMessage))
            else:
                resp = None
        except EOFError:
            log.debug("done polling")
            return {'FINISHED'}

        if resp is not None:
            if resp.exception is not None:
                raise resp.exception
            elif resp.data is not None:
                self.handle_response(resp.data) #TODO: make this a customizable callback
            # this should allow chaining of multiple subprocess in a single operator

    decoratify(cls, 'execute', decorate_execute)
    decoratify(cls, 'invoke', decorate_invoke)
    setattr(cls, 'poll_subprocess', poll_subprocess)

    return cls

def subprocess_function(func, *args, pipe, **kwargs):
    """
    Wrapper which feeds the return val of `func` into the latter end of a pipe
    """
    def wrapper(*args, **kwargs):
        # we have to explicitly close the end of the pipe we are NOT using,
        # otherwise no exception will be generated when the other process closes its end.
        pipe[0].close()

        try:
            return_val = func(*args, **kwargs)
            pipe[1].send(SubprocessMessage(data=return_val))
        except Exception as err:
            log.debug("Caught exception from subprocess: %s", err)
            pipe[1].send(SubprocessMessage(exception=err))
            raise
        finally:
            pipe[1].close()

    return wrapper

class SubprocessMessage:
    """
    Container for communications between child processes and the parent process
    """
    #TODO: currently only used for child -> parent, might need something different for parent -> child?
    def __init__(self, exception=None, data=None):
        self.exception = exception
        self.data = data