# Released under the MIT License. See LICENSE for details. # """Functionality for sending and responding to messages. Supports static typing for message types and possible return types. """ from __future__ import annotations from typing import TYPE_CHECKING import traceback import logging import json from efro.error import CleanError from efro.dataclassio import (is_ioprepped_dataclass, dataclass_to_dict, dataclass_from_dict) from efro.message._message import (Message, Response, ErrorResponse, EmptyResponse, UnregisteredMessageIDError) if TYPE_CHECKING: from typing import Any, Literal class MessageProtocol: """Wrangles a set of message types, formats, and response types. Both endpoints must be using a compatible Protocol for communication to succeed. To maintain Protocol compatibility between revisions, all message types must retain the same id, message attr storage names must not change, newly added attrs must have default values, etc. """ def __init__(self, message_types: dict[int, type[Message]], response_types: dict[int, type[Response]], preserve_clean_errors: bool = True, receiver_logs_exceptions: bool = True, receiver_returns_stack_traces: bool = False) -> None: """Create a protocol with a given configuration. Note that common response types are automatically registered with (unchanging negative ids) so they don't need to be passed explicitly (but can be if a different id is desired). If 'preserve_clean_errors' is True, efro.error.CleanError exceptions raised on the receiver end will result in a matching CleanError raised back on the sender. All other Exception types come across as efro.error.RemoteError. When 'receiver_logs_exceptions' is True, any uncaught Exceptions on the receiver end will be logged there via logging.exception() (in addition to the usual behavior of returning an ErrorResponse to the sender). This is good to leave enabled if your intention is to never return ErrorResponses. Looser setups making routine use of CleanErrors or whatnot may want to disable this, however. If 'receiver_returns_stack_traces' is True, stringified stack traces will be returned to the sender for exceptions occurring on the receiver end. This can make debugging easier but should only be used when the client is trusted to see such info. """ self.message_types_by_id: dict[int, type[Message]] = {} self.message_ids_by_type: dict[type[Message], int] = {} self.response_types_by_id: dict[int, type[Response]] = {} self.response_ids_by_type: dict[type[Response], int] = {} for m_id, m_type in message_types.items(): # Make sure only valid message types were passed and each # id was assigned only once. assert isinstance(m_id, int) assert m_id >= 0 assert (is_ioprepped_dataclass(m_type) and issubclass(m_type, Message)) assert self.message_types_by_id.get(m_id) is None self.message_types_by_id[m_id] = m_type self.message_ids_by_type[m_type] = m_id for r_id, r_type in response_types.items(): assert isinstance(r_id, int) assert r_id >= 0 assert (is_ioprepped_dataclass(r_type) and issubclass(r_type, Response)) assert self.response_types_by_id.get(r_id) is None self.response_types_by_id[r_id] = r_type self.response_ids_by_type[r_type] = r_id # Go ahead and auto-register a few common response types # if the user has not done so explicitly. Use unique negative # IDs which will never change or overlap with user ids. def _reg_if_not(reg_tp: type[Response], reg_id: int) -> None: if reg_tp in self.response_ids_by_type: return assert self.response_types_by_id.get(reg_id) is None self.response_types_by_id[reg_id] = reg_tp self.response_ids_by_type[reg_tp] = reg_id _reg_if_not(ErrorResponse, -1) _reg_if_not(EmptyResponse, -2) # Some extra-thorough validation in debug mode. if __debug__: # Make sure all Message types' return types are valid # and have been assigned an ID as well. all_response_types: set[type[Response]] = set() for m_id, m_type in message_types.items(): m_rtypes = m_type.get_response_types() assert isinstance(m_rtypes, list) assert m_rtypes, ( f'Message type {m_type} specifies no return types.') assert len(set(m_rtypes)) == len(m_rtypes) # check dups all_response_types.update(m_rtypes) for cls in all_response_types: assert is_ioprepped_dataclass(cls) assert issubclass(cls, Response) if cls not in self.response_ids_by_type: raise ValueError( f'Possible response type {cls} needs to be included' f' in response_types for this protocol.') # Make sure all registered types have unique base names. # We can take advantage of this to generate cleaner looking # protocol modules. Can revisit if this is ever a problem. mtypenames = set(tp.__name__ for tp in self.message_ids_by_type) if len(mtypenames) != len(message_types): raise ValueError( 'message_types contains duplicate __name__s;' ' all types are required to have unique names.') self.preserve_clean_errors = preserve_clean_errors self.receiver_logs_exceptions = receiver_logs_exceptions self.receiver_returns_stack_traces = receiver_returns_stack_traces @staticmethod def encode_dict(obj: dict) -> str: """Json-encode a provided dict.""" return json.dumps(obj, separators=(',', ':')) def message_to_dict(self, message: Message) -> dict: """Encode a message to a json ready dict.""" return self._to_dict(message, self.message_ids_by_type, 'message') def response_to_dict(self, response: Response) -> dict: """Encode a response to a json ready dict.""" return self._to_dict(response, self.response_ids_by_type, 'response') def error_to_response(self, exc: Exception) -> Response: """Translate an error to a response.""" # Log any errors we got during handling if so desired. if self.receiver_logs_exceptions: logging.exception('Error handling message.') # If anything goes wrong, return a ErrorResponse instead. if isinstance(exc, CleanError) and self.preserve_clean_errors: return ErrorResponse(error_message=str(exc), error_type=ErrorResponse.ErrorType.CLEAN) return ErrorResponse( error_message=(traceback.format_exc() if self.receiver_returns_stack_traces else 'An internal error has occurred.'), error_type=ErrorResponse.ErrorType.OTHER) def _to_dict(self, message: Any, ids_by_type: dict[type, int], opname: str) -> dict: """Encode a message to a json string for transport.""" m_id: int | None = ids_by_type.get(type(message)) if m_id is None: raise TypeError(f'{opname} type is not registered in protocol:' f' {type(message)}') out = {'t': m_id, 'm': dataclass_to_dict(message)} return out @staticmethod def decode_dict(data: str) -> dict: """Decode data to a dict.""" out = json.loads(data) assert isinstance(out, dict) return out def message_from_dict(self, data: dict) -> Message: """Decode a message from a json string.""" out = self._from_dict(data, self.message_types_by_id, 'message') assert isinstance(out, Message) return out def response_from_dict(self, data: dict) -> Response: """Decode a response from a json string.""" out = self._from_dict(data, self.response_types_by_id, 'response') assert isinstance(out, Response) return out # Weeeird; we get mypy errors returning dict[int, type] but # dict[int, typing.Type] or dict[int, type[Any]] works.. def _from_dict(self, data: dict, types_by_id: dict[int, type[Any]], opname: str) -> Any: """Decode a message from a json string.""" msgdict: dict | None m_id = data.get('t') # Allow omitting 'm' dict if its empty. msgdict = data.get('m', {}) assert isinstance(m_id, int) assert isinstance(msgdict, dict) # Decode this particular type. msgtype = types_by_id.get(m_id) if msgtype is None: raise UnregisteredMessageIDError( f'Got unregistered {opname} id of {m_id}.') return dataclass_from_dict(msgtype, msgdict) def _get_module_header(self, part: Literal['sender', 'receiver'], extra_import_code: str | None = None) -> str: """Return common parts of generated modules.""" # pylint: disable=too-many-locals, too-many-branches import textwrap tpimports: dict[str, list[str]] = {} imports: dict[str, list[str]] = {} single_message_type = len(self.message_ids_by_type) == 1 msgtypes = list(self.message_ids_by_type) if part == 'sender': msgtypes.append(Message) for msgtype in msgtypes: tpimports.setdefault(msgtype.__module__, []).append(msgtype.__name__) rsptypes = list(self.response_ids_by_type) if part == 'sender': rsptypes.append(Response) for rsp_tp in rsptypes: # Skip these as they don't actually show up in code. if rsp_tp is EmptyResponse or rsp_tp is ErrorResponse: continue if (single_message_type and part == 'sender' and rsp_tp is not Response): # We need to cast to the single supported response type # in this case so need response types at runtime. imports.setdefault(rsp_tp.__module__, []).append(rsp_tp.__name__) else: tpimports.setdefault(rsp_tp.__module__, []).append(rsp_tp.__name__) import_lines = '' tpimport_lines = '' for module, names in sorted(imports.items()): jnames = ', '.join(names) line = f'from {module} import {jnames}' if len(line) > 79: # Recreate in a wrapping-friendly form. line = f'from {module} import ({jnames})' import_lines += f'{line}\n' for module, names in sorted(tpimports.items()): jnames = ', '.join(names) line = f'from {module} import {jnames}' if len(line) > 75: # Account for indent # Recreate in a wrapping-friendly form. line = f'from {module} import ({jnames})' tpimport_lines += f'{line}\n' if part == 'sender': import_lines += ('from efro.message import MessageSender,' ' BoundMessageSender') tpimport_typing_extras = '' else: if single_message_type: import_lines += ('from efro.message import (MessageReceiver,' ' BoundMessageReceiver, Message, Response)') else: import_lines += ('from efro.message import MessageReceiver,' ' BoundMessageReceiver') tpimport_typing_extras = ', Awaitable' if extra_import_code is not None: import_lines += f'\n{extra_import_code}\n' ovld = ', overload' if not single_message_type else '' tpimport_lines = textwrap.indent(tpimport_lines, ' ') baseimps = ['Any'] if part == 'receiver': baseimps.append('Callable') baseimps_s = ', '.join(baseimps) out = ('# Released under the MIT License. See LICENSE for details.\n' f'#\n' f'"""Auto-generated {part} module. Do not edit by hand."""\n' f'\n' f'from __future__ import annotations\n' f'\n' f'from typing import TYPE_CHECKING{ovld}\n' f'\n' f'{import_lines}\n' f'\n' f'if TYPE_CHECKING:\n' f' from typing import {baseimps_s}' f'{tpimport_typing_extras}\n' f'{tpimport_lines}' f'\n' f'\n') return out def do_create_sender_module( self, basename: str, protocol_create_code: str, enable_sync_sends: bool, enable_async_sends: bool, private: bool = False, protocol_module_level_import_code: str | None = None) -> str: """Used by create_sender_module(); do not call directly.""" # pylint: disable=too-many-locals import textwrap msgtypes = list(self.message_ids_by_type.keys()) ppre = '_' if private else '' out = self._get_module_header( 'sender', extra_import_code=protocol_module_level_import_code) ccind = textwrap.indent(protocol_create_code, ' ') out += (f'class {ppre}{basename}(MessageSender):\n' f' """Protocol-specific sender."""\n' f'\n' f' def __init__(self) -> None:\n' f'{ccind}\n' f' super().__init__(protocol)\n' f'\n' f' def __get__(self,\n' f' obj: Any,\n' f' type_in: Any = None)' f' -> {ppre}Bound{basename}:\n' f' return {ppre}Bound{basename}' f'(obj, self)\n' f'\n' f'\n' f'class {ppre}Bound{basename}(BoundMessageSender):\n' f' """Protocol-specific bound sender."""\n') def _filt_tp_name(rtype: type[Response]) -> str: # We accept None to equal EmptyResponse so reflect that # in the type annotation. return 'None' if rtype is EmptyResponse else rtype.__name__ # Define handler() overloads for all registered message types. if msgtypes: for async_pass in False, True: if async_pass and not enable_async_sends: continue if not async_pass and not enable_sync_sends: continue pfx = 'async ' if async_pass else '' sfx = '_async' if async_pass else '' awt = 'await ' if async_pass else '' how = 'asynchronously' if async_pass else 'synchronously' if len(msgtypes) == 1: # Special case: with a single message types we don't # use overloads. msgtype = msgtypes[0] msgtypevar = msgtype.__name__ rtypes = msgtype.get_response_types() if len(rtypes) > 1: rtypevar = ' | '.join(_filt_tp_name(t) for t in rtypes) else: rtypevar = _filt_tp_name(rtypes[0]) out += (f'\n' f' {pfx}def send{sfx}(self,' f' message: {msgtypevar})' f' -> {rtypevar}:\n' f' """Send a message {how}."""\n' f' out = {awt}self._sender.' f'send{sfx}(self._obj, message)\n' f' assert isinstance(out, {rtypevar})\n' f' return out\n') else: for msgtype in msgtypes: msgtypevar = msgtype.__name__ rtypes = msgtype.get_response_types() if len(rtypes) > 1: rtypevar = ' | '.join( _filt_tp_name(t) for t in rtypes) else: rtypevar = _filt_tp_name(rtypes[0]) out += (f'\n' f' @overload\n' f' {pfx}def send{sfx}(self,' f' message: {msgtypevar})' f' -> {rtypevar}:\n' f' ...\n') out += (f'\n' f' {pfx}def send{sfx}(self, message: Message)' f' -> Response | None:\n' f' """Send a message {how}."""\n' f' return {awt}self._sender.' f'send{sfx}(self._obj, message)\n') return out def do_create_receiver_module( self, basename: str, protocol_create_code: str, is_async: bool, private: bool = False, protocol_module_level_import_code: str | None = None) -> str: """Used by create_receiver_module(); do not call directly.""" # pylint: disable=too-many-locals import textwrap desc = 'asynchronous' if is_async else 'synchronous' ppre = '_' if private else '' msgtypes = list(self.message_ids_by_type.keys()) out = self._get_module_header( 'receiver', extra_import_code=protocol_module_level_import_code) ccind = textwrap.indent(protocol_create_code, ' ') out += (f'class {ppre}{basename}(MessageReceiver):\n' f' """Protocol-specific {desc} receiver."""\n' f'\n' f' is_async = {is_async}\n' f'\n' f' def __init__(self) -> None:\n' f'{ccind}\n' f' super().__init__(protocol)\n' f'\n' f' def __get__(\n' f' self,\n' f' obj: Any,\n' f' type_in: Any = None,\n' f' ) -> {ppre}Bound{basename}:\n' f' return {ppre}Bound{basename}(' f'obj, self)\n') # Define handler() overloads for all registered message types. def _filt_tp_name(rtype: type[Response]) -> str: # We accept None to equal EmptyResponse so reflect that # in the type annotation. return 'None' if rtype is EmptyResponse else rtype.__name__ if msgtypes: cbgn = 'Awaitable[' if is_async else '' cend = ']' if is_async else '' if len(msgtypes) == 1: # Special case: when we have a single message type we don't # use overloads. msgtype = msgtypes[0] msgtypevar = msgtype.__name__ rtypes = msgtype.get_response_types() if len(rtypes) > 1: rtypevar = ' | '.join(_filt_tp_name(t) for t in rtypes) else: rtypevar = _filt_tp_name(rtypes[0]) rtypevar = f'{cbgn}{rtypevar}{cend}' out += ( f'\n' f' def handler(\n' f' self,\n' f' call: Callable[[Any, {msgtypevar}], ' f'{rtypevar}],\n' f' )' f' -> Callable[[Any, {msgtypevar}], {rtypevar}]:\n' f' """Decorator to register message handlers."""\n' f' from typing import cast, Callable, Any\n' f' self.register_handler(cast(Callable' f'[[Any, Message], Response], call))\n' f' return call\n') else: for msgtype in msgtypes: msgtypevar = msgtype.__name__ rtypes = msgtype.get_response_types() if len(rtypes) > 1: rtypevar = ' | '.join(_filt_tp_name(t) for t in rtypes) else: rtypevar = _filt_tp_name(rtypes[0]) rtypevar = f'{cbgn}{rtypevar}{cend}' out += (f'\n' f' @overload\n' f' def handler(\n' f' self,\n' f' call: Callable[[Any, {msgtypevar}], ' f'{rtypevar}],\n' f' )' f' -> Callable[[Any, {msgtypevar}], {rtypevar}]:\n' f' ...\n') out += ( '\n' ' def handler(self, call: Callable) -> Callable:\n' ' """Decorator to register message handlers."""\n' ' self.register_handler(call)\n' ' return call\n') out += (f'\n' f'\n' f'class {ppre}Bound{basename}(BoundMessageReceiver):\n' f' """Protocol-specific bound receiver."""\n') if is_async: out += ( '\n' ' async def handle_raw_message(self,\n' ' message: str,\n' ' raise_unregistered: bool =' ' False) -> str:\n' ' """Asynchronously handle a raw incoming message."""\n' ' return await self._receiver.handle_raw_message_async(' '\n' ' self._obj, message, raise_unregistered)\n') else: out += ( '\n' ' def handle_raw_message(self,\n' ' message: str,\n' ' raise_unregistered: bool = False)' ' -> str:\n' ' """Synchronously handle a raw incoming message."""\n' ' return self._receiver.handle_raw_message(' 'self._obj, message,\n' ' ' 'raise_unregistered)\n') return out