forked from goodboy/tractor
1
0
Fork 0

Prepare to offer (dynamic) `.msg.Codec` overrides

By simply allowing an input `codec: tuple` of funcs for now to the
`MsgpackTCPStream` transport but, ideally wrapping this in a `Codec`
type with an API for dynamic extension of the interchange lib's msg
processing settings. Right now we're tied to `msgspec.msgpack` for this
transport but with the right design this can likely extend to other libs
in the future.

Relates to starting feature work toward #36, #196, #365.
old_msg_types
Tyler Goodlet 2024-03-25 16:31:16 -04:00
parent 72b4dc1461
commit 496dce57a8
1 changed files with 38 additions and 5 deletions

View File

@ -30,6 +30,7 @@ import struct
import typing import typing
from typing import ( from typing import (
Any, Any,
Callable,
runtime_checkable, runtime_checkable,
Protocol, Protocol,
Type, Type,
@ -123,6 +124,16 @@ class MsgpackTCPStream(MsgTransport):
stream: trio.SocketStream, stream: trio.SocketStream,
prefix_size: int = 4, prefix_size: int = 4,
# XXX optionally provided codec pair for `msgspec`:
# https://jcristharif.com/msgspec/extending.html#mapping-to-from-native-types
#
# TODO: define this as a `Codec` struct which can be
# overriden dynamically by the application/runtime.
codec: tuple[
Callable[[Any], Any]|None, # coder
Callable[[type, Any], Any]|None, # decoder
]|None = None,
) -> None: ) -> None:
self.stream = stream self.stream = stream
@ -138,12 +149,18 @@ class MsgpackTCPStream(MsgTransport):
# public i guess? # public i guess?
self.drained: list[dict] = [] self.drained: list[dict] = []
self.recv_stream = BufferedReceiveStream(transport_stream=stream) self.recv_stream = BufferedReceiveStream(
transport_stream=stream
)
self.prefix_size = prefix_size self.prefix_size = prefix_size
# TODO: struct aware messaging coders # TODO: struct aware messaging coders
self.encode = msgspec.msgpack.Encoder().encode self.encode = msgspec.msgpack.Encoder(
self.decode = msgspec.msgpack.Decoder().decode # dict[str, Any]) enc_hook=codec[0] if codec else None,
).encode
self.decode = msgspec.msgpack.Decoder(
dec_hook=codec[1] if codec else None,
).decode
async def _iter_packets(self) -> AsyncGenerator[dict, None]: async def _iter_packets(self) -> AsyncGenerator[dict, None]:
'''Yield packets from the underlying stream. '''Yield packets from the underlying stream.
@ -349,9 +366,25 @@ class Channel:
stream: trio.SocketStream, stream: trio.SocketStream,
type_key: tuple[str, str]|None = None, type_key: tuple[str, str]|None = None,
# XXX optionally provided codec pair for `msgspec`:
# https://jcristharif.com/msgspec/extending.html#mapping-to-from-native-types
codec: tuple[
Callable[[Any], Any], # coder
Callable[[type, Any], Any], # decoder
]|None = None,
) -> MsgTransport: ) -> MsgTransport:
type_key = type_key or self._transport_key type_key = (
self._transport = get_msg_transport(type_key)(stream) type_key
or
self._transport_key
)
self._transport = get_msg_transport(
type_key
)(
stream,
codec=codec,
)
return self._transport return self._transport
def __repr__(self) -> str: def __repr__(self) -> str: