sync changes with master

dist/ba_data/python/efro/dataclassio/_api.py

@@ -18,7 +18,7 @@ from efro.dataclassio._inputter import _Inputter
 from efro.dataclassio._base import Codec
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any
 
 T = TypeVar('T')
 
@@ -69,7 +69,7 @@ def dataclass_to_dict(obj: Any,
 def dataclass_to_json(obj: Any,
                       coerce_to_float: bool = True,
                       pretty: bool = False,
-                      sort_keys: Optional[bool] = None) -> str:
+                      sort_keys: bool | None = None) -> str:
     """Utility function; return a json string from a dataclass instance.
 
     Basically json.dumps(dataclass_to_dict(...)).

dist/ba_data/python/efro/dataclassio/_base.py

@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, get_args
 from typing import _AnnotatedAlias  # type: ignore
 
 if TYPE_CHECKING:
-    from typing import Any, Optional, Callable, Union
+    from typing import Any, Callable
 
 # Types which we can pass through as-is.
 SIMPLE_TYPES = {int, bool, str, float, type(None)}
@@ -31,8 +31,7 @@ def _raise_type_error(fieldpath: str, valuetype: type,
     if len(expected) == 1:
         expected_str = expected[0].__name__
     else:
-        names = ', '.join(t.__name__ for t in expected)
-        expected_str = f'Union[{names}]'
+        expected_str = ' | '.join(t.__name__ for t in expected)
     raise TypeError(f'Invalid value type for "{fieldpath}";'
                     f' expected "{expected_str}", got'
                     f' "{valuetype.__name__}".')
@@ -128,21 +127,21 @@ class IOAttrs:
 
     MISSING = _MissingType()
 
-    storagename: Optional[str] = None
+    storagename: str | None = None
     store_default: bool = True
     whole_days: bool = False
     whole_hours: bool = False
     soft_default: Any = MISSING
-    soft_default_factory: Union[Callable[[], Any], _MissingType] = MISSING
+    soft_default_factory: Callable[[], Any] | _MissingType = MISSING
 
     def __init__(
         self,
-        storagename: Optional[str] = storagename,
+        storagename: str | None = storagename,
         store_default: bool = store_default,
         whole_days: bool = whole_days,
         whole_hours: bool = whole_hours,
         soft_default: Any = MISSING,
-        soft_default_factory: Union[Callable[[], Any], _MissingType] = MISSING,
+        soft_default_factory: Callable[[], Any] | _MissingType = MISSING,
     ):
 
         # Only store values that differ from class defaults to keep
@@ -216,12 +215,12 @@ def _get_origin(anntype: Any) -> Any:
     return anntype if origin is None else origin
 
 
-def _parse_annotated(anntype: Any) -> tuple[Any, Optional[IOAttrs]]:
+def _parse_annotated(anntype: Any) -> tuple[Any, IOAttrs | None]:
     """Parse Annotated() constructs, returning annotated type & IOAttrs."""
     # If we get an Annotated[foo, bar, eep] we take
     # foo as the actual type, and we look for IOAttrs instances in
     # bar/eep to affect our behavior.
-    ioattrs: Optional[IOAttrs] = None
+    ioattrs: IOAttrs | None = None
     if isinstance(anntype, _AnnotatedAlias):
         annargs = get_args(anntype)
         for annarg in annargs[1:]:

dist/ba_data/python/efro/dataclassio/_inputter.py

@@ -23,7 +23,7 @@ from efro.dataclassio._base import (Codec, _parse_annotated, EXTRA_ATTRS_ATTR,
 from efro.dataclassio._prep import PrepSession
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any
 
     from efro.dataclassio._base import IOAttrs
     from efro.dataclassio._outputter import _Outputter
@@ -44,7 +44,7 @@ class _Inputter(Generic[T]):
         self._coerce_to_float = coerce_to_float
         self._allow_unknown_attrs = allow_unknown_attrs
         self._discard_unknown_attrs = discard_unknown_attrs
-        self._soft_default_validator: Optional[_Outputter] = None
+        self._soft_default_validator: _Outputter | None = None
 
         if not allow_unknown_attrs and discard_unknown_attrs:
             raise ValueError('discard_unknown_attrs cannot be True'
@@ -63,7 +63,7 @@ class _Inputter(Generic[T]):
         return out
 
     def _value_from_input(self, cls: type, fieldpath: str, anntype: Any,
-                          value: Any, ioattrs: Optional[IOAttrs]) -> Any:
+                          value: Any, ioattrs: IOAttrs | None) -> Any:
         """Convert an assigned value to what a dataclass field expects."""
         # pylint: disable=too-many-return-statements
         # pylint: disable=too-many-branches
@@ -270,7 +270,7 @@ class _Inputter(Generic[T]):
                                                         fieldpath=fieldpath)
 
     def _dict_from_input(self, cls: type, fieldpath: str, anntype: Any,
-                         value: Any, ioattrs: Optional[IOAttrs]) -> Any:
+                         value: Any, ioattrs: IOAttrs | None) -> Any:
         # pylint: disable=too-many-branches
         # pylint: disable=too-many-locals
 
@@ -370,7 +370,7 @@ class _Inputter(Generic[T]):
 
     def _sequence_from_input(self, cls: type, fieldpath: str, anntype: Any,
                              value: Any, seqtype: type,
-                             ioattrs: Optional[IOAttrs]) -> Any:
+                             ioattrs: IOAttrs | None) -> Any:
 
         # Because we are json-centric, we expect a list for all sequences.
         if type(value) is not list:
@@ -396,7 +396,7 @@ class _Inputter(Generic[T]):
             for i in value)
 
     def _datetime_from_input(self, cls: type, fieldpath: str, value: Any,
-                             ioattrs: Optional[IOAttrs]) -> Any:
+                             ioattrs: IOAttrs | None) -> Any:
 
         # For firestore we expect a datetime object.
         if self._codec is Codec.FIRESTORE:
@@ -428,7 +428,7 @@ class _Inputter(Generic[T]):
         return out
 
     def _tuple_from_input(self, cls: type, fieldpath: str, anntype: Any,
-                          value: Any, ioattrs: Optional[IOAttrs]) -> Any:
+                          value: Any, ioattrs: IOAttrs | None) -> Any:
 
         out: list = []
 

dist/ba_data/python/efro/dataclassio/_outputter.py

@@ -23,7 +23,7 @@ from efro.dataclassio._base import (Codec, _parse_annotated, EXTRA_ATTRS_ATTR,
 from efro.dataclassio._prep import PrepSession
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any
     from efro.dataclassio._base import IOAttrs
 
 
@@ -64,7 +64,7 @@ class _Outputter:
                                                           recursion_level=0)
         assert prep is not None
         fields = dataclasses.fields(obj)
-        out: Optional[dict[str, Any]] = {} if self._create else None
+        out: dict[str, Any] | None = {} if self._create else None
         for field in fields:
             fieldname = field.name
             if fieldpath:
@@ -118,15 +118,16 @@ class _Outputter:
         if isinstance(extra_attrs, dict):
             if not _is_valid_for_codec(extra_attrs, self._codec):
                 raise TypeError(
-                    f'Extra attrs on {fieldpath} contains data type(s)'
-                    f' not supported by json.')
+                    f'Extra attrs on \'{fieldpath}\' contains data type(s)'
+                    f' not supported by \'{self._codec.value}\' codec:'
+                    f' {extra_attrs}.')
             if self._create:
                 assert out is not None
                 out.update(extra_attrs)
         return out
 
     def _process_value(self, cls: type, fieldpath: str, anntype: Any,
-                       value: Any, ioattrs: Optional[IOAttrs]) -> Any:
+                       value: Any, ioattrs: IOAttrs | None) -> Any:
         # pylint: disable=too-many-return-statements
         # pylint: disable=too-many-branches
         # pylint: disable=too-many-statements
@@ -307,7 +308,7 @@ class _Outputter:
         return value
 
     def _process_dict(self, cls: type, fieldpath: str, anntype: Any,
-                      value: dict, ioattrs: Optional[IOAttrs]) -> Any:
+                      value: dict, ioattrs: IOAttrs | None) -> Any:
         # pylint: disable=too-many-branches
         if not isinstance(value, dict):
             raise TypeError(f'Expected a dict for {fieldpath};'
@@ -329,7 +330,7 @@ class _Outputter:
 
         # Ok; we've got a definite key type (which we verified as valid
         # during prep). Make sure all keys match it.
-        out: Optional[dict] = {} if self._create else None
+        out: dict | None = {} if self._create else None
         keyanntype, valanntype = childtypes
 
         # str keys we just export directly since that's supported by json.

dist/ba_data/python/efro/dataclassio/_prep.py

@@ -21,7 +21,7 @@ from efro.dataclassio._base import (_parse_annotated, _get_origin,
                                     SIMPLE_TYPES)
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any
     from efro.dataclassio._base import IOAttrs
 
 T = TypeVar('T')
@@ -115,12 +115,12 @@ class PrepData:
 class PrepSession:
     """Context for a prep."""
 
-    def __init__(self, explicit: bool, globalns: Optional[dict] = None):
+    def __init__(self, explicit: bool, globalns: dict | None = None):
         self.explicit = explicit
         self.globalns = globalns
 
     def prep_dataclass(self, cls: type,
-                       recursion_level: int) -> Optional[PrepData]:
+                       recursion_level: int) -> PrepData | None:
         """Run prep on a dataclass if necessary and return its prep data.
 
         The only case where this will return None is for recursive types
@@ -232,7 +232,7 @@ class PrepSession:
         return prepdata
 
     def prep_type(self, cls: type, attrname: str, anntype: Any,
-                  ioattrs: Optional[IOAttrs], recursion_level: int) -> None:
+                  ioattrs: IOAttrs | None, recursion_level: int) -> None:
         """Run prep on a dataclass."""
         # pylint: disable=too-many-return-statements
         # pylint: disable=too-many-branches

dist/ba_data/python/efro/message/_message.py

@@ -42,12 +42,6 @@ class Response:
 # Some standard response types:
 
 
-class ErrorType(Enum):
-    """Type of error that occurred in remote message handling."""
-    OTHER = 0
-    CLEAN = 1
-
-
 @ioprepped
 @dataclass
 class ErrorResponse(Response):
@@ -56,6 +50,13 @@ class ErrorResponse(Response):
     This type is unique in that it is not returned to the user; it
     instead results in a local exception being raised.
     """
+
+    class ErrorType(Enum):
+        """Type of error that occurred in remote message handling."""
+        OTHER = 0
+        CLEAN = 1
+        LOCAL = 2
+
     error_message: Annotated[str, IOAttrs('m')]
     error_type: Annotated[ErrorType, IOAttrs('e')] = ErrorType.OTHER
 

dist/ba_data/python/efro/message/_module.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING
 from efro.message._protocol import MessageProtocol
 
 if TYPE_CHECKING:
-    from typing import Optional
+    pass
 
 
 def create_sender_module(
@@ -20,8 +20,8 @@ def create_sender_module(
     enable_sync_sends: bool,
     enable_async_sends: bool,
     private: bool = False,
-    protocol_module_level_import_code: Optional[str] = None,
-    build_time_protocol_create_code: Optional[str] = None,
+    protocol_module_level_import_code: str | None = None,
+    build_time_protocol_create_code: str | None = None,
 ) -> str:
     """Create a Python module defining a MessageSender subclass.
 
@@ -59,8 +59,8 @@ def create_receiver_module(
     protocol_create_code: str,
     is_async: bool,
     private: bool = False,
-    protocol_module_level_import_code: Optional[str] = None,
-    build_time_protocol_create_code: Optional[str] = None,
+    protocol_module_level_import_code: str | None = None,
+    build_time_protocol_create_code: str | None = None,
 ) -> str:
     """"Create a Python module defining a MessageReceiver subclass.
 

dist/ba_data/python/efro/message/_protocol.py

@@ -15,8 +15,7 @@ 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, ErrorType,
-                                   UnregisteredMessageIDError)
+                                   EmptyResponse, UnregisteredMessageIDError)
 
 if TYPE_CHECKING:
     from typing import Any, Literal
@@ -141,11 +140,11 @@ class MessageProtocol:
         # 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=ErrorType.CLEAN)
+                                 error_type=ErrorResponse.ErrorType.CLEAN)
         return ErrorResponse(
             error_message=(traceback.format_exc() if self.trusted_sender else
                            'An unknown error has occurred.'),
-            error_type=ErrorType.OTHER)
+            error_type=ErrorResponse.ErrorType.OTHER)
 
     def _to_dict(self, message: Any, ids_by_type: dict[type, int],
                  opname: str) -> dict:

dist/ba_data/python/efro/message/_receiver.py

@@ -15,7 +15,7 @@ from efro.message._message import (Message, Response, EmptyResponse,
                                    ErrorResponse, UnregisteredMessageIDError)
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional, Union
+    from typing import Any, Callable, Awaitable
 
     from efro.message._protocol import MessageProtocol
 
@@ -50,14 +50,19 @@ class MessageReceiver:
     def __init__(self, protocol: MessageProtocol) -> None:
         self.protocol = protocol
         self._handlers: dict[type[Message], Callable] = {}
-        self._decode_filter_call: Optional[Callable[[Any, dict, Message],
-                                                    None]] = None
-        self._encode_filter_call: Optional[Callable[[Any, Response, dict],
-                                                    None]] = None
+        self._decode_filter_call: Callable[[Any, dict, Message],
+                                           None] | None = None
+        self._encode_filter_call: Callable[
+            [Any, Message | None, Response, dict], None] | None = None
+
+        # TODO: don't currently have async encode equivalent
+        # or either for sender; can add as needed.
+        self._decode_filter_async_call: Callable[[Any, dict, Message],
+                                                 Awaitable[None]] | None = None
 
     # noinspection PyProtectedMember
     def register_handler(
-            self, call: Callable[[Any, Message], Optional[Response]]) -> None:
+            self, call: Callable[[Any, Message], Response | None]) -> None:
         """Register a handler call.
 
         The message type handled by the call is determined by its
@@ -101,7 +106,7 @@ class MessageReceiver:
         assert issubclass(msgtype, Message)
 
         ret = anns.get('return')
-        responsetypes: tuple[Union[type[Any], type[None]], ...]
+        responsetypes: tuple[type[Any] | type[None], ...]
 
         # Return types can be a single type or a union of types.
         if isinstance(ret, (_GenericAlias, types.UnionType)):
@@ -152,15 +157,29 @@ class MessageReceiver:
         """Function decorator for defining a decode filter.
 
         Decode filters can be used to extract extra data from incoming
-        message dicts.
+        message dicts. This version will work for both handle_raw_message()
+        and handle_raw_message_async()
         """
         assert self._decode_filter_call is None
         self._decode_filter_call = call
         return call
 
+    def decode_filter_async_method(
+        self, call: Callable[[Any, dict, Message], Awaitable[None]]
+    ) -> Callable[[Any, dict, Message], Awaitable[None]]:
+        """Function decorator for defining a decode filter.
+
+        Decode filters can be used to extract extra data from incoming
+        message dicts. Note that this version will only work with
+        handle_raw_message_async().
+        """
+        assert self._decode_filter_async_call is None
+        self._decode_filter_async_call = call
+        return call
+
     def encode_filter_method(
-        self, call: Callable[[Any, Response, dict], None]
-    ) -> Callable[[Any, Response, dict], None]:
+        self, call: Callable[[Any, Message | None, Response, dict], None]
+    ) -> Callable[[Any, Message | None, Response, dict], None]:
         """Function decorator for defining an encode filter.
 
         Encode filters can be used to add extra data to the message
@@ -183,21 +202,38 @@ class MessageReceiver:
                 else:
                     raise TypeError(msg)
 
-    def _decode_incoming_message(self, bound_obj: Any,
-                                 msg: str) -> tuple[Message, type[Message]]:
+    def _decode_incoming_message_base(self, bound_obj: Any,
+                                      msg: str) -> tuple[Any, dict, Message]:
         # Decode the incoming message.
         msg_dict = self.protocol.decode_dict(msg)
         msg_decoded = self.protocol.message_from_dict(msg_dict)
-        msgtype = type(msg_decoded)
-        assert issubclass(msgtype, Message)
+        assert isinstance(msg_decoded, Message)
         if self._decode_filter_call is not None:
             self._decode_filter_call(bound_obj, msg_dict, msg_decoded)
+        return bound_obj, msg_dict, msg_decoded
 
-        return msg_decoded, msgtype
+    def _decode_incoming_message(self, bound_obj: Any, msg: str) -> Message:
+        bound_obj, _msg_dict, msg_decoded = (
+            self._decode_incoming_message_base(bound_obj=bound_obj, msg=msg))
 
-    def encode_user_response(self, bound_obj: Any,
-                             response: Optional[Response],
-                             msgtype: type[Message]) -> str:
+        # If they've set an async filter but are calling sync
+        # handle_raw_message() its likely a bug.
+        assert self._decode_filter_async_call is None
+
+        return msg_decoded
+
+    async def _decode_incoming_message_async(self, bound_obj: Any,
+                                             msg: str) -> Message:
+        bound_obj, msg_dict, msg_decoded = (self._decode_incoming_message_base(
+            bound_obj=bound_obj, msg=msg))
+
+        if self._decode_filter_async_call is not None:
+            await self._decode_filter_async_call(bound_obj, msg_dict,
+                                                 msg_decoded)
+        return msg_decoded
+
+    def encode_user_response(self, bound_obj: Any, message: Message,
+                             response: Response | None) -> str:
         """Encode a response provided by the user for sending."""
 
         # A return value of None equals EmptyResponse.
@@ -207,18 +243,21 @@ class MessageReceiver:
         assert isinstance(response, Response)
         # (user should never explicitly return error-responses)
         assert not isinstance(response, ErrorResponse)
-        assert type(response) in msgtype.get_response_types()
+        assert type(response) in message.get_response_types()
         response_dict = self.protocol.response_to_dict(response)
         if self._encode_filter_call is not None:
-            self._encode_filter_call(bound_obj, response, response_dict)
+            self._encode_filter_call(bound_obj, message, response,
+                                     response_dict)
         return self.protocol.encode_dict(response_dict)
 
-    def encode_error_response(self, bound_obj: Any, exc: Exception) -> str:
+    def encode_error_response(self, bound_obj: Any, message: Message | None,
+                              exc: Exception) -> str:
         """Given an error, return a response ready for sending."""
         response = self.protocol.error_to_response(exc)
         response_dict = self.protocol.response_to_dict(response)
         if self._encode_filter_call is not None:
-            self._encode_filter_call(bound_obj, response, response_dict)
+            self._encode_filter_call(bound_obj, message, response,
+                                     response_dict)
         return self.protocol.encode_dict(response_dict)
 
     def handle_raw_message(self,
@@ -233,21 +272,22 @@ class MessageReceiver:
         error responses returned to the sender.
         """
         assert not self.is_async, "can't call sync handler on async receiver"
+        msg_decoded: Message | None = None
         try:
-            msg_decoded, msgtype = self._decode_incoming_message(
-                bound_obj, msg)
+            msg_decoded = self._decode_incoming_message(bound_obj, msg)
+            msgtype = type(msg_decoded)
             handler = self._handlers.get(msgtype)
             if handler is None:
                 raise RuntimeError(f'Got unhandled message type: {msgtype}.')
             response = handler(bound_obj, msg_decoded)
             assert isinstance(response, (Response, type(None)))
-            return self.encode_user_response(bound_obj, response, msgtype)
+            return self.encode_user_response(bound_obj, msg_decoded, response)
 
         except Exception as exc:
             if (raise_unregistered
                     and isinstance(exc, UnregisteredMessageIDError)):
                 raise
-            return self.encode_error_response(bound_obj, exc)
+            return self.encode_error_response(bound_obj, msg_decoded, exc)
 
     async def handle_raw_message_async(
             self,
@@ -259,21 +299,23 @@ class MessageReceiver:
         The return value is the raw response to the message.
         """
         assert self.is_async, "can't call async handler on sync receiver"
+        msg_decoded: Message | None = None
         try:
-            msg_decoded, msgtype = self._decode_incoming_message(
+            msg_decoded = await self._decode_incoming_message_async(
                 bound_obj, msg)
+            msgtype = type(msg_decoded)
             handler = self._handlers.get(msgtype)
             if handler is None:
                 raise RuntimeError(f'Got unhandled message type: {msgtype}.')
             response = await handler(bound_obj, msg_decoded)
             assert isinstance(response, (Response, type(None)))
-            return self.encode_user_response(bound_obj, response, msgtype)
+            return self.encode_user_response(bound_obj, msg_decoded, response)
 
         except Exception as exc:
             if (raise_unregistered
                     and isinstance(exc, UnregisteredMessageIDError)):
                 raise
-            return self.encode_error_response(bound_obj, exc)
+            return self.encode_error_response(bound_obj, msg_decoded, exc)
 
 
 class BoundMessageReceiver:
@@ -294,5 +336,12 @@ class BoundMessageReceiver:
         return self._receiver.protocol
 
     def encode_error_response(self, exc: Exception) -> str:
-        """Given an error, return a response ready to send."""
-        return self._receiver.encode_error_response(self._obj, exc)
+        """Given an error, return a response ready to send.
+
+        This should be used for any errors that happen outside of
+        of standard handle_raw_message calls. Any errors within those
+        calls should be automatically returned as encoded strings.
+        """
+        # Passing None for Message here; we would only have that available
+        # for things going wrong in the handler (which this is not for).
+        return self._receiver.encode_error_response(self._obj, None, exc)

dist/ba_data/python/efro/message/_sender.py

@@ -6,13 +6,14 @@ Supports static typing for message types and possible return types.
 
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING, TypeVar
 
 from efro.error import CleanError, RemoteError
-from efro.message._message import (EmptyResponse, ErrorResponse, ErrorType)
+from efro.message._message import EmptyResponse, ErrorResponse
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional, Awaitable
+    from typing import Any, Callable, Awaitable
 
     from efro.message._message import Message, Response
     from efro.message._protocol import MessageProtocol
@@ -42,13 +43,13 @@ class MessageSender:
 
     def __init__(self, protocol: MessageProtocol) -> None:
         self.protocol = protocol
-        self._send_raw_message_call: Optional[Callable[[Any, str], str]] = None
-        self._send_async_raw_message_call: Optional[Callable[
-            [Any, str], Awaitable[str]]] = None
-        self._encode_filter_call: Optional[Callable[[Any, Message, dict],
-                                                    None]] = None
-        self._decode_filter_call: Optional[Callable[[Any, dict, Response],
-                                                    None]] = None
+        self._send_raw_message_call: Callable[[Any, str], str] | None = None
+        self._send_async_raw_message_call: Callable[
+            [Any, str], Awaitable[str]] | None = None
+        self._encode_filter_call: Callable[[Any, Message, dict],
+                                           None] | None = None
+        self._decode_filter_call: Callable[[Any, Message, dict, Response],
+                                           None] | None = None
 
     def send_method(
             self, call: Callable[[Any, str],
@@ -79,8 +80,8 @@ class MessageSender:
         return call
 
     def decode_filter_method(
-        self, call: Callable[[Any, dict, Response], None]
-    ) -> Callable[[Any, dict, Response], None]:
+        self, call: Callable[[Any, Message, dict, Response], None]
+    ) -> Callable[[Any, Message, dict, Response], None]:
         """Function decorator for defining a decode filter.
 
         Decode filters can be used to extract extra data from incoming
@@ -90,71 +91,137 @@ class MessageSender:
         self._decode_filter_call = call
         return call
 
-    def send(self, bound_obj: Any, message: Message) -> Optional[Response]:
-        """Send a message and receive a response.
+    def send(self, bound_obj: Any, message: Message) -> Response | None:
+        """Send a message synchronously."""
+        return self.send_split_part_2(
+            message=message,
+            raw_response=self.send_split_part_1(
+                bound_obj=bound_obj,
+                message=message,
+            ),
+        )
 
-        Will encode the message for transport and call dispatch_raw_message()
+    async def send_async(self, bound_obj: Any,
+                         message: Message) -> Response | None:
+        """Send a message asynchronously."""
+        return self.send_split_part_2(
+            message=message,
+            raw_response=await self.send_split_part_1_async(
+                bound_obj=bound_obj,
+                message=message,
+            ),
+        )
+
+    def send_split_part_1(self, bound_obj: Any, message: Message) -> Response:
+        """Send a message synchronously.
+
+        Generally you can just call send(); these split versions are
+        for when message sending and response handling need to happen
+        in different contexts/threads.
         """
         if self._send_raw_message_call is None:
             raise RuntimeError('send() is unimplemented for this type.')
 
-        msg_encoded = self.encode_message(bound_obj, message)
-
+        msg_encoded = self._encode_message(bound_obj, message)
         response_encoded = self._send_raw_message_call(bound_obj, msg_encoded)
+        return self._decode_raw_response(bound_obj, message, response_encoded)
 
-        response = self.decode_response(bound_obj, response_encoded)
+    async def send_split_part_1_async(self, bound_obj: Any,
+                                      message: Message) -> Response:
+        """Send a message asynchronously.
+
+        Generally you can just call send(); these split versions are
+        for when message sending and response handling need to happen
+        in different contexts/threads.
+        """
+
+        if self._send_async_raw_message_call is None:
+            raise RuntimeError('send_async() is unimplemented for this type.')
+
+        msg_encoded = self._encode_message(bound_obj, message)
+        response_encoded = await self._send_async_raw_message_call(
+            bound_obj, msg_encoded)
+
+        return self._decode_raw_response(bound_obj, message, response_encoded)
+
+    def send_split_part_2(self, message: Message,
+                          raw_response: Response) -> Response | None:
+        """Complete message sending (both sync and async).
+
+        Generally you can just call send(); these split versions are
+        for when message sending and response handling need to happen
+        in different contexts/threads.
+        """
+        response = self._unpack_raw_response(raw_response)
         assert (response is None
                 or type(response) in type(message).get_response_types())
         return response
 
-    def encode_message(self, bound_obj: Any, message: Message) -> str:
+    def _encode_message(self, bound_obj: Any, message: Message) -> str:
         """Encode a message for sending."""
         msg_dict = self.protocol.message_to_dict(message)
         if self._encode_filter_call is not None:
             self._encode_filter_call(bound_obj, message, msg_dict)
         return self.protocol.encode_dict(msg_dict)
 
-    def decode_response(self, bound_obj: Any,
-                        response_encoded: str) -> Optional[Response]:
-        """Decode, filter, and possibly act on raw response data."""
-        response_dict = self.protocol.decode_dict(response_encoded)
-        response = self.protocol.response_from_dict(response_dict)
-        if self._decode_filter_call is not None:
-            self._decode_filter_call(bound_obj, response_dict, response)
+    def _decode_raw_response(self, bound_obj: Any, message: Message,
+                             response_encoded: str) -> Response:
+        """Create a Response from returned data.
 
-        # Special case: if we get EmptyResponse, we simply return None.
-        if isinstance(response, EmptyResponse):
+        These Responses may encapsulate things like remote errors and
+        should not be handed directly to users. _unpack_raw_response()
+        should be used to translate to special values like None or raise
+        Exceptions. This function itself should never raise Exceptions.
+        """
+        try:
+            response_dict = self.protocol.decode_dict(response_encoded)
+            response = self.protocol.response_from_dict(response_dict)
+            if self._decode_filter_call is not None:
+                self._decode_filter_call(bound_obj, message, response_dict,
+                                         response)
+        except Exception:
+            # If we got to this point, we successfully communicated
+            # with the other end so errors represent protocol mismatches
+            # or other invalid data. For now let's just log it but perhaps
+            # we'd want to somehow embed it in the ErrorResponse to be raised
+            # directly to the user later.
+            logging.exception('Error decoding raw response')
+            response = ErrorResponse(
+                error_message=
+                'Error decoding raw response; see log for details.',
+                error_type=ErrorResponse.ErrorType.LOCAL)
+        return response
+
+    def _unpack_raw_response(self, raw_response: Response) -> Response | None:
+        """Given a raw Response, unpacks to special values or Exceptions.
+
+        The result of this call is what should be passed to users.
+        For complex messaging situations such as response callbacks
+        operating across different threads, this last stage should be
+        run such that any raised Exception is active when the callback
+        fires; not on the thread where the message was sent.
+        """
+        # EmptyResponse translates to None
+        if isinstance(raw_response, EmptyResponse):
             return None
 
-        # Special case: a remote error occurred. Raise a local Exception
-        # instead of returning the message.
-        if isinstance(response, ErrorResponse):
-            if (self.protocol.preserve_clean_errors
-                    and response.error_type is ErrorType.CLEAN):
-                raise CleanError(response.error_message)
-            raise RemoteError(response.error_message)
+        # Some error occurred. Raise a local Exception for it.
+        if isinstance(raw_response, ErrorResponse):
 
-        return response
+            # If something went wrong on our end of the connection,
+            # don't say it was a remote error.
+            if raw_response.error_type is ErrorResponse.ErrorType.LOCAL:
+                raise RuntimeError(raw_response.error_message)
 
-    async def send_async(self, bound_obj: Any,
-                         message: Message) -> Optional[Response]:
-        """Send a message asynchronously using asyncio.
+            # If they want to support clean errors, do those.
+            if (self.protocol.preserve_clean_errors and
+                    raw_response.error_type is ErrorResponse.ErrorType.CLEAN):
+                raise CleanError(raw_response.error_message)
 
-        The message will be encoded for transport and passed to
-        dispatch_raw_message_async.
-        """
-        if self._send_async_raw_message_call is None:
-            raise RuntimeError('send_async() is unimplemented for this type.')
+            # In all other cases, just say something went wrong 'out there'.
+            raise RemoteError(raw_response.error_message)
 
-        msg_encoded = self.encode_message(bound_obj, message)
-
-        response_encoded = await self._send_async_raw_message_call(
-            bound_obj, msg_encoded)
-
-        response = self.decode_response(bound_obj, response_encoded)
-        assert (response is None
-                or type(response) in type(message).get_response_types())
-        return response
+        return raw_response
 
 
 class BoundMessageSender:
@@ -171,20 +238,34 @@ class BoundMessageSender:
         """Protocol associated with this sender."""
         return self._sender.protocol
 
-    def send_untyped(self, message: Message) -> Optional[Response]:
+    def send_untyped(self, message: Message) -> Response | None:
         """Send a message synchronously.
 
         Whenever possible, use the send() call provided by generated
         subclasses instead of this; it will provide better type safety.
         """
         assert self._obj is not None
-        return self._sender.send(self._obj, message)
+        return self._sender.send(bound_obj=self._obj, message=message)
 
-    async def send_async_untyped(self, message: Message) -> Optional[Response]:
+    async def send_async_untyped(self, message: Message) -> Response | None:
         """Send a message asynchronously.
 
         Whenever possible, use the send_async() call provided by generated
         subclasses instead of this; it will provide better type safety.
         """
         assert self._obj is not None
-        return await self._sender.send_async(self._obj, message)
+        return await self._sender.send_async(bound_obj=self._obj,
+                                             message=message)
+
+    async def send_split_part_1_async_untyped(self,
+                                              message: Message) -> Response:
+        """Split send (part 1 of 2)."""
+        assert self._obj is not None
+        return await self._sender.send_split_part_1_async(bound_obj=self._obj,
+                                                          message=message)
+
+    def send_split_part_2_untyped(self, message: Message,
+                                  raw_response: Response) -> Response | None:
+        """Split send (part 2 of 2)."""
+        return self._sender.send_split_part_2(message=message,
+                                              raw_response=raw_response)

dist/ba_data/python/efro/rpc.py

@@ -19,7 +19,7 @@ from efro.dataclassio import (dataclass_to_json, dataclass_from_json,
                               ioprepped, IOAttrs)
 
 if TYPE_CHECKING:
-    from typing import Literal, Awaitable, Callable, Optional
+    from typing import Literal, Awaitable, Callable
 
 # Terminology:
 # Packet: A chunk of data consisting of a type and some type-dependent
@@ -33,6 +33,8 @@ class _PacketType(Enum):
     KEEPALIVE = 1
     MESSAGE = 2
     RESPONSE = 3
+    MESSAGE_BIG = 4
+    RESPONSE_BIG = 5
 
 
 _BYTE_ORDER: Literal['big'] = 'big'
@@ -49,14 +51,20 @@ class _PeerInfo:
     keepalive_interval: Annotated[float, IOAttrs('k')]
 
 
-OUR_PROTOCOL = 1
+# Note: we are expected to be forward and backward compatible; we can
+# increment protocol freely and expect everyone else to still talk to us.
+# Likewise we should retain logic to communicate with older protocols.
+# Protocol history:
+# 1 - initial release
+# 2 - gained big (32-bit len val) package/response packets
+OUR_PROTOCOL = 2
 
 
 class _InFlightMessage:
     """Represents a message that is out on the wire."""
 
     def __init__(self) -> None:
-        self._response: Optional[bytes] = None
+        self._response: bytes | None = None
         self._got_response = asyncio.Event()
         self.wait_task = asyncio.create_task(self._wait())
 
@@ -126,7 +134,7 @@ class RPCEndpoint:
         self._out_packets: list[bytes] = []
         self._have_out_packets = asyncio.Event()
         self._run_called = False
-        self._peer_info: Optional[_PeerInfo] = None
+        self._peer_info: _PeerInfo | None = None
         self._keepalive_interval = keepalive_interval
         self._keepalive_timeout = keepalive_timeout
 
@@ -135,7 +143,7 @@ class RPCEndpoint:
         self._tasks: list[weakref.ref[asyncio.Task]] = []
 
         # When we last got a keepalive or equivalent (time.monotonic value)
-        self._last_keepalive_receive_time: Optional[float] = None
+        self._last_keepalive_receive_time: float | None = None
 
         # (Start near the end to make sure our looping logic is sound).
         self._next_message_id = 65530
@@ -193,7 +201,7 @@ class RPCEndpoint:
 
     async def send_message(self,
                            message: bytes,
-                           timeout: Optional[float] = None) -> bytes:
+                           timeout: float | None = None) -> bytes:
         """Send a message to the peer and return a response.
 
         If timeout is not provided, the default will be used.
@@ -201,21 +209,38 @@ class RPCEndpoint:
         for any reason.
         """
         self._check_env()
-        if len(message) > 65535:
-            raise RuntimeError('Message cannot be larger than 65535 bytes')
 
         if self._closing:
             raise CommunicationError('Endpoint is closed')
 
-        # Go with 16 bit looping value for message_id.
+        # We need to know their protocol, so if we haven't gotten a handshake
+        # from them yet, just wait.
+        while self._peer_info is None:
+            await asyncio.sleep(0.01)
+        assert self._peer_info is not None
+
+        if self._peer_info.protocol == 1:
+            if len(message) > 65535:
+                raise RuntimeError('Message cannot be larger than 65535 bytes')
+
+        # message_id is a 16 bit looping value.
         message_id = self._next_message_id
         self._next_message_id = (self._next_message_id + 1) % 65536
 
-        # Payload consists of type (1b), message_id (2b), len (2b), and data.
-        self._enqueue_outgoing_packet(
-            _PacketType.MESSAGE.value.to_bytes(1, _BYTE_ORDER) +
-            message_id.to_bytes(2, _BYTE_ORDER) +
-            len(message).to_bytes(2, _BYTE_ORDER) + message)
+        if len(message) > 65535:
+            # Payload consists of type (1b), message_id (2b),
+            # len (4b), and data.
+            self._enqueue_outgoing_packet(
+                _PacketType.MESSAGE_BIG.value.to_bytes(1, _BYTE_ORDER) +
+                message_id.to_bytes(2, _BYTE_ORDER) +
+                len(message).to_bytes(4, _BYTE_ORDER) + message)
+        else:
+            # Payload consists of type (1b), message_id (2b),
+            # len (2b), and data.
+            self._enqueue_outgoing_packet(
+                _PacketType.MESSAGE.value.to_bytes(1, _BYTE_ORDER) +
+                message_id.to_bytes(2, _BYTE_ORDER) +
+                len(message).to_bytes(2, _BYTE_ORDER) + message)
 
         # Make an entry so we know this message is out there.
         assert message_id not in self._in_flight_messages
@@ -381,17 +406,27 @@ class RPCEndpoint:
                 self._last_keepalive_receive_time = time.monotonic()
 
             elif mtype is _PacketType.MESSAGE:
-                await self._handle_message_packet()
+                await self._handle_message_packet(big=False)
+
+            elif mtype is _PacketType.MESSAGE_BIG:
+                await self._handle_message_packet(big=True)
 
             elif mtype is _PacketType.RESPONSE:
-                await self._handle_response_packet()
+                await self._handle_response_packet(big=False)
+
+            elif mtype is _PacketType.RESPONSE_BIG:
+                await self._handle_response_packet(big=True)
 
             else:
                 assert_never(mtype)
 
-    async def _handle_message_packet(self) -> None:
+    async def _handle_message_packet(self, big: bool) -> None:
+        assert self._peer_info is not None
         msgid = await self._read_int_16()
-        msglen = await self._read_int_16()
+        if big:
+            msglen = await self._read_int_32()
+        else:
+            msglen = await self._read_int_16()
         msg = await self._reader.readexactly(msglen)
         if self._debug_print_io:
             self._debug_print_call(f'{self._label}: received message {msgid}'
@@ -408,9 +443,14 @@ class RPCEndpoint:
         self._debug_print_call(
             f'{self._label}: done handling message at {self._tm()}.')
 
-    async def _handle_response_packet(self) -> None:
+    async def _handle_response_packet(self, big: bool) -> None:
+        assert self._peer_info is not None
         msgid = await self._read_int_16()
-        rsplen = await self._read_int_16()
+        # Protocol 2 gained 32 bit data lengths.
+        if big:
+            rsplen = await self._read_int_32()
+        else:
+            rsplen = await self._read_int_16()
         if self._debug_print_io:
             self._debug_print_call(f'{self._label}: received response {msgid}'
                                    f' of size {rsplen} at {self._tm()}.')
@@ -520,12 +560,25 @@ class RPCEndpoint:
             logging.exception('Error handling message')
             return
 
+        assert self._peer_info is not None
+
+        if self._peer_info.protocol == 1:
+            if len(response) > 65535:
+                raise RuntimeError(
+                    'Response cannot be larger than 65535 bytes')
+
         # Now send back our response.
         # Payload consists of type (1b), msgid (2b), len (2b), and data.
-        self._enqueue_outgoing_packet(
-            _PacketType.RESPONSE.value.to_bytes(1, _BYTE_ORDER) +
-            message_id.to_bytes(2, _BYTE_ORDER) +
-            len(response).to_bytes(2, _BYTE_ORDER) + response)
+        if len(response) > 65535:
+            self._enqueue_outgoing_packet(
+                _PacketType.RESPONSE_BIG.value.to_bytes(1, _BYTE_ORDER) +
+                message_id.to_bytes(2, _BYTE_ORDER) +
+                len(response).to_bytes(4, _BYTE_ORDER) + response)
+        else:
+            self._enqueue_outgoing_packet(
+                _PacketType.RESPONSE.value.to_bytes(1, _BYTE_ORDER) +
+                message_id.to_bytes(2, _BYTE_ORDER) +
+                len(response).to_bytes(2, _BYTE_ORDER) + response)
 
     async def _read_int_8(self) -> int:
         return int.from_bytes(await self._reader.readexactly(1), _BYTE_ORDER)