modelverse/kernel/modelverse_kernel/request_handler.py

import sys
import modelverse_kernel.primitives as primitive_functions
import modelverse_kernel.jit as jit
from collections import defaultdict

class RequestHandler(object):
    def __init__(self):
        self.generator_stack = []
        self.handlers = {"CALL": self.execute_call,
                         "CALL_ARGS": self.execute_call_args,
                         "CALL_KWARGS": self.execute_call_kwargs,
                         "SLEEP": self.execute_sleep}

    def push_generator(self, gen):
        self.generator_stack.append(gen)

    def handle_request(self, reply):
        while self.generator_stack:
            try:
                gen = self.generator_stack[-1]
                requests = gen.send(reply)

                # Generated new request, so process
                if requests and requests[0][0] in self.handlers:
                        # Known request, so process that one
                        if len(requests) > 1:
                            raise Exception("CALL_* and SLEEP operations MUST be split in individual yields")

                        # Try to add a new generator to branch into
                        reply = None
                        self.generator_stack.append(None)
                        # This next command potentially raises a finished message already, meaning that we should stop already
                        # We avoid an extra try/except block by putting the None on the stack already
                        self.generator_stack[-1] = self.handlers[requests[0][0]](requests[0][1])
                else:
                    # MvS request, so forward that instead
                    return requests

            except StopIteration:
                # Exception, so finished execution of this generator, passing on None to the caller
                del self.generator_stack[-1]
                reply = [None]
                
            except primitive_functions.PrimitiveFinished as ex:
                # Exception, so finished execution of this generator, passing on ex.result to the caller
                del self.generator_stack[-1]
                reply = [ex.result]

            except primitive_functions.SleepKernel:
                # Processing sleep, so pop its generator and reraise
                del self.generator_stack[-1]
                raise

    def execute_call(self, request_args):
        return request_args[0]

    def execute_call_kwargs(self, request_args):
        return request_args[0](**(request_args[1]))

    def execute_call_args(self, request_args):
        return request_args[0](*(request_args[1]))

    def execute_sleep(self, request_args):
        raise primitive_functions.SleepKernel(request_args[0], request_args[1])