Source code for wsrpc_aiohttp.websocket.common

import abc
import asyncio
import logging
import types
from collections import defaultdict
from functools import partial
import typing as t

import aiohttp

from . import decorators
from .abc import (
    Proxy, AbstactWSRPC, FrameMappingItemType, RouteType, EventListenerType
)
from .route import Route
from .tools import Singleton, awaitable, loads
from wsrpc_aiohttp.signal import Signal


[docs]class WSRPCError(Exception): pass
[docs]class ClientException(WSRPCError): __slots__ = ("type", "message", "raw") def __init__(self, payload): self.type = payload.get("type") self.message = payload.get("message") self.raw = payload
class PingTimeoutError(WSRPCError): pass def ping(_, **kwargs): return kwargs log = logging.getLogger(__name__) class Nothing(Singleton): def __repr__(self): return self.__class__.__name__ CallItem = t.NamedTuple( "CallItem", ( ("serial", t.Optional[int]), ("method", t.Union[Nothing, str, None]), ("error", t.Union[Nothing, t.Any]), ("result", t.Union[Nothing, t.Any]), ("params", t.Optional[t.Union[t.List, t.Dict]]), ), ) RouteCollectionType = t.DefaultDict[ t.Type[AbstactWSRPC], t.Dict[str, RouteType] ] ClientCollectionType = t.DefaultDict[ t.Type[AbstactWSRPC], t.Dict[str, AbstactWSRPC] ] LocksCollectionType = t.DefaultDict[int, asyncio.Lock] FutureCollectionType = t.DefaultDict[int, asyncio.Future] EventListenerCollectionType = t.Set[EventListenerType] def _route_maker() -> t.Dict[str, RouteType]: return {"ping": ping} # type: ignore
[docs]class WSRPCBase(AbstactWSRPC): """ Common WSRPC abstraction """ _ROUTES = defaultdict(_route_maker) # type: RouteCollectionType _CLIENTS = defaultdict(dict) # type: ClientCollectionType _CLEAN_LOCK_TIMEOUT = 2 # type: t.Union[int, float] __slots__ = ( "_handlers", "_loop", "_pending_tasks", "_locks", "_futures", "_serial", "_timeout", "_event_listeners", "_message_type_mapping", ) ON_CALL_START = Signal() ON_CALL_SUCCESS = Signal() ON_CALL_FAIL = Signal() def __init__(self, loop: asyncio.AbstractEventLoop = None, timeout: t.Union[int, float] = None): self._loop = loop or asyncio.get_event_loop() self._handlers = {} # type: t.Dict[str, RouteType] self._pending_tasks = set() # type: t.Set[asyncio.Task] self._serial = 0 self._timeout = timeout self._locks = defaultdict(asyncio.Lock) # type: LocksCollectionType self._futures = defaultdict( self._loop.create_future ) # type: FutureCollectionType self._event_listeners = set() # type: EventListenerCollectionType self._message_type_mapping = self._create_type_mapping() def _create_type_mapping(self) -> FrameMappingItemType: return types.MappingProxyType( { aiohttp.WSMsgType.TEXT: self.handle_message, aiohttp.WSMsgType.BINARY: self.handle_binary, aiohttp.WSMsgType.CLOSE: self.close, aiohttp.WSMsgType.CLOSED: self.close, } ) def _create_task(self, coro): task = self._loop.create_task(coro) # type: asyncio.Task self._pending_tasks.add(task) task.add_done_callback(partial(self._pending_tasks.remove)) return task def _call_later(self, timer, callback, *args, **kwargs): def handler(): self._create_task(awaitable(callback)(*args, **kwargs)) self._pending_tasks.add(self._loop.call_later(timer, handler))
[docs] async def close(self, message=None): """ Cancel all pending tasks """ async def task_waiter(task): if not (hasattr(task, "__iter__") or hasattr(task, "__aiter__")): return try: await task except asyncio.CancelledError: pass except Exception: log.exception( "Unhandled exception when closing client connection" ) if message: log.info("Closing WebSocket because message %r received", message) for task in tuple(self._pending_tasks): task.cancel() if hasattr(task, "cancelled") and not task.cancelled(): self._loop.create_task(task_waiter(task))
async def handle_binary(self, message: aiohttp.WSMessage): log.warning("Unhandled message %r %r", message.type, message.data) async def _call_method(self, call_item: CallItem): try: if not isinstance(call_item.method, Nothing) and call_item.serial: log.debug( "Acquiring lock for %r serial %r", self, call_item.serial ) async with self._locks[call_item.serial]: args, kwargs = self.prepare_args(call_item.params) return await self.handle_method( call_item.method, call_item.serial, args, kwargs ) elif not isinstance(call_item.result, Nothing): return await self.handle_result( call_item.serial, call_item.result ) elif not isinstance(call_item.error, Nothing): return await self.handle_error( call_item.serial, call_item.error ) else: return await self.handle_result(call_item.serial, None) except Exception as e: log.exception(e) if call_item.serial: await self._send( error=self._format_error(e), id=call_item.serial ) finally: self._call_later( self._CLEAN_LOCK_TIMEOUT, self.__clean_lock, call_item.serial, ) @staticmethod def _parse_message(data: dict) -> CallItem: message_id = data.get("id") # type: t.Optional[int] if message_id and not isinstance(message_id, int): raise ValueError message_method = data.get( "method", Nothing() ) # type: t.Union[str, Nothing, None] message_result = data.get( "result", Nothing() ) # type: t.Union[str, Nothing, None] message_error = data.get( "error", Nothing() ) # type: t.Union[str, Nothing, None] message_params = data.get( "params", None ) # type: t.Union[t.List[t.Any], t.Dict[t.Any, t.Any], None] return CallItem( serial=message_id, method=message_method, result=message_result, error=message_error, params=message_params, ) async def handle_message(self, message: aiohttp.WSMessage): # noinspection PyTypeChecker, PyNoneFunctionAssignment data = message.json(loads=loads) # type: dict log.debug("Got message: %r", data) serial = data.get("id") if serial is None: return await self.handle_event(data) call_item = self._parse_message(data) await self._call_method(call_item) async def _on_message(self, msg: aiohttp.WSMessage): async def unknown_method(msg: aiohttp.WSMessage): log.warning("Unhandled message %r %r", msg.type, msg.data) handler = self._message_type_mapping.get(msg.type, unknown_method) self._create_task(awaitable(handler)(msg)) @classmethod def get_routes(cls) -> t.Dict[str, RouteType]: return cls._ROUTES[cls] @classmethod def get_clients(cls) -> t.Dict[str, AbstactWSRPC]: return cls._CLIENTS[cls] @property def routes(self) -> t.Dict[str, RouteType]: """ Property which contains the socket routes """ return self.get_routes() @property def clients(self) -> t.Dict[str, AbstactWSRPC]: """ Property which contains the socket clients """ return self.get_clients() @staticmethod def _prepare_args(args): arguments = [] kwargs = {} if isinstance(args, type(None)): return arguments, kwargs if isinstance(args, list): arguments.extend(args) elif isinstance(args, dict): kwargs.update(args) else: arguments.append(args) return arguments, kwargs def prepare_args(self, args): return self._prepare_args(args) @staticmethod def is_route(func): return hasattr(func, "__self__") and isinstance(func.__self__, Route) async def handle_method(self, method, serial, args, kwargs): await self.ON_CALL_START.call( method=method, serial=serial, args=args, kwargs=kwargs, ) callee = self.resolver(method) if not self.is_route(callee): a = [self] a.extend(args) args = a func = partial(callee, *args, **kwargs) try: result = await self._executor(func) except Exception as err: await self.ON_CALL_FAIL.call( method=method, serial=serial, args=args, kwargs=kwargs, err=err, ) raise await self.ON_CALL_SUCCESS.call( method=method, serial=serial, args=args, kwargs=kwargs, result=result, ) await self._send(result=result, id=serial) async def handle_result(self, serial, result): cb = self._futures.pop(serial, None) if not cb or cb.done(): return cb.set_result(result) async def handle_error(self, serial, error): self._reject(serial, error) log.error("Client return error: \n\t%r", error) def __clean_lock(self, serial): if serial not in self._locks: return log.debug("Release and delete lock for %s serial %s", self, serial) self._locks.pop(serial) async def handle_event(self, event): for listener in self._event_listeners: self._loop.call_soon(listener, event) @abc.abstractmethod async def _send(self, **kwargs): raise NotImplementedError @staticmethod def _format_error(e): return {"type": str(type(e).__name__), "message": str(e)} def _reject(self, serial, error): future = self._futures.get(serial) if not future: return future.set_exception(ClientException(error)) def _unresolvable(self, func_name, *args, **kwargs): raise NotImplementedError( 'Callback function "%r" not implemented' % func_name ) def resolver(self, func_name): class_name, method = ( func_name.split(".", 1) if "." in func_name else (func_name, "init") ) callee = self.routes.get(class_name, self._unresolvable) if isinstance(callee, decorators.ProxyBase): callee = callee.func condition = ( callee == self._unresolvable or isinstance(getattr(callee, "__self__", None), Route) or ( not isinstance(callee, (types.FunctionType, types.MethodType)) and issubclass(callee, Route) ) ) if condition: if class_name not in self._handlers: self._handlers[class_name] = callee(self) return self._handlers[class_name](method) callee = self.routes.get(func_name, self._unresolvable) if hasattr(callee, "__call__"): return callee else: raise NotImplementedError( "Method call of {0} is not implemented".format(repr(callee)) ) def _get_serial(self): self._serial += 2 return self._serial
[docs] async def call(self, func: str, timeout=None, **kwargs): """ Method for call remote function Remote methods allows only kwargs as arguments. You might use functions as route or classes .. code-block:: python async def remote_function(socket: WSRPCBase, *, foo, bar): # call function from the client-side await self.socket.proxy.ping() return foo + bar class RemoteClass(WebSocketRoute): # this method executes when remote side call route name asyc def init(self): # call function from the client-side await self.socket.proxy.ping() async def make_something(self, foo, bar): return foo + bar """ serial = self._get_serial() future = self._futures[serial] payload = dict(id=serial, method=func, params=kwargs) log.info( 'Sending request #%r "%s(%r)" to the client.', serial, func, kwargs, ) await self._send(**payload) result = await asyncio.wait_for( future, timeout=timeout or self._timeout ) return result
async def emit(self, event): await self._send(**event)
[docs] @classmethod def add_route(cls, route: str, handler: RouteType): """ Expose local function through RPC :param route: Name which function will be aliased for this function. Remote side should call function by this name. :param handler: Function or Route class (classes based on :class:`wsrpc_aiohttp.WebSocketRoute`). For route classes the public methods will be registered automatically. .. note:: Route classes might be initialized only once for the each socket instance. In case the method of class will be called first, :func:`wsrpc_aiohttp.WebSocketRoute.init` will be called without params before callable method. """ assert callable(handler) or isinstance(handler, Route) if callable(handler): handler = decorators.proxy(handler) cls.get_routes()[route] = handler
def add_event_listener(self, func: t.Callable[[dict], t.Any]): self._event_listeners.add(func) def remove_event_listeners(self, func): return self._event_listeners.remove(func)
[docs] @classmethod def remove_route(cls, route: str, fail=True): """ Removes route by name. If `fail=True` an exception will be raised in case the route was not found. """ if fail: cls.get_routes().pop(route) else: cls.get_routes().pop(route, None)
def __repr__(self): if hasattr(self, "id"): return "<RPCWebSocket: ID[{0}]>".format(self.id) else: return "<RPCWebsocket: {0} (waiting)>".format(self.__hash__()) @abc.abstractmethod async def _executor(self, func): raise NotImplementedError @property def proxy(self): """ Special property which allow run the remote functions by `dot` notation .. code-block:: python # calls remote function with name ping await client.proxy.ping() # full equivalent of await client.call('ping') """ return Proxy(self.call)
__all__ = ("Route", "WSRPCBase", "ClientException", "WSRPCError")