forked from goodboy/tractor
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
parent
72b4dc1461
commit
496dce57a8
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue