from functools import wraps from ssb.packet_stream import PSMessageType class MuxRPCRequestHandler(object): def __init__(self, ps_handler): self.ps_handler = ps_handler def __await__(self): return self.ps_handler.__await__() class MuxRPCSourceHandler(object): def __init__(self, ps_handler): self.ps_handler = ps_handler async def __aiter__(self): async for msg in self.ps_handler: yield msg class MuxRPCSinkHandlerMixin(object): def send(self, msg, msg_type=PSMessageType.JSON): self.connection.send(msg, stream=True, msg_type=msg_type, req=self.req) class MuxRPCDuplexHandler(MuxRPCSinkHandlerMixin, MuxRPCSourceHandler): def __init__(self, ps_handler, connection, req): super(MuxRPCDuplexHandler, self).__init__(ps_handler) self.connection = connection self.req = req class MuxRPCSinkHandler(MuxRPCSinkHandlerMixin): def __init__(self, connection, req): self.connection = connection self.req = req def _get_appropriate_api_handler(type_, connection, ps_handler, req): if type_ in {'sync', 'async'}: return MuxRPCRequestHandler(ps_handler) elif type_ == 'source': return MuxRPCSourceHandler(ps_handler) elif type_ == 'sink': return MuxRPCSinkHandler(connection, req) elif type_ == 'duplex': return MuxRPCDuplexHandler(ps_handler, connection, req) class MuxRPCAPIException(Exception): pass class MuxRPCRequest(object): @classmethod def from_message(cls, message): body = message.body return cls('.'.join(body['name']), body['args']) def __init__(self, name, args): self.name = name self.args = args def __repr__(self): return ''.format(self) class MuxRPCMessage(object): @classmethod def from_message(cls, message): return cls(message.body) def __init__(self, body): self.body = body def __repr__(self): return ''.format(self) class MuxRPCAPI(object): def __init__(self): self.handlers = {} self.connection = None async def __await__(self): async for req_message in self.connection: body = req_message.body if req_message is None: return if isinstance(body, dict) and body.get('name'): self.process(self.connection, MuxRPCRequest.from_message(req_message)) def add_connection(self, connection): self.connection = connection def define(self, name): def _handle(f): self.handlers[name] = f @wraps(f) def _f(*args, **kwargs): return f(*args, **kwargs) return f return _handle def process(self, connection, request): handler = self.handlers.get(request.name) if not handler: raise MuxRPCAPIException('Method {} not found!'.format(request.name)) handler(connection, request) def call(self, name, args, type_='sync'): old_counter = self.connection.req_counter ps_handler = self.connection.send({ 'name': name.split('.'), 'args': args, 'type': type_ }, stream=type_ in {'sink', 'source', 'duplex'}) return _get_appropriate_api_handler(type_, self.connection, ps_handler, old_counter)