Add custom `MsgCodec.__repr__()`
Sure makes console grokability a lot better by showing only the customizeable fields. Further, clean up `mk_codec()` a bunch by removing the `ipc_msg_spec` param since we don't plan to support another msg-set (for now) which allows cleaning out a buncha logic that was mostly just a source of bugs.. Also, - add temporary `log.info()` around codec application. - throw in some sanity `assert`s to `limit_msg_spec()`. - add but mask out the `extend_msg_spec()` idea since it seems `msgspec` won't allow `Decoder.type` extensions when using a custom `dec_hook()` for some extension type.. (not sure what approach to take here yet).runtime_to_msgspec
parent
dbc445ff9d
commit
322e015d32
|
@ -37,6 +37,7 @@ from contextlib import (
|
||||||
# ContextVar,
|
# ContextVar,
|
||||||
# Token,
|
# Token,
|
||||||
# )
|
# )
|
||||||
|
import textwrap
|
||||||
from typing import (
|
from typing import (
|
||||||
Any,
|
Any,
|
||||||
Callable,
|
Callable,
|
||||||
|
@ -59,7 +60,9 @@ from tractor.msg.types import (
|
||||||
mk_msg_spec,
|
mk_msg_spec,
|
||||||
MsgType,
|
MsgType,
|
||||||
)
|
)
|
||||||
|
from tractor.log import get_logger
|
||||||
|
|
||||||
|
log = get_logger(__name__)
|
||||||
|
|
||||||
# TODO: overall IPC msg-spec features (i.e. in this mod)!
|
# TODO: overall IPC msg-spec features (i.e. in this mod)!
|
||||||
#
|
#
|
||||||
|
@ -87,6 +90,27 @@ class MsgCodec(Struct):
|
||||||
|
|
||||||
pld_spec: Union[Type[Struct]]|None
|
pld_spec: Union[Type[Struct]]|None
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
speclines: str = textwrap.indent(
|
||||||
|
self.pformat_msg_spec(),
|
||||||
|
prefix=' '*3,
|
||||||
|
)
|
||||||
|
body: str = textwrap.indent(
|
||||||
|
f'|_lib = {self.lib.__name__!r}\n'
|
||||||
|
f'|_enc_hook: {self.enc.enc_hook}\n'
|
||||||
|
f'|_dec_hook: {self.dec.dec_hook}\n'
|
||||||
|
f'|_pld_spec: {self.pld_spec_str}\n'
|
||||||
|
# f'|\n'
|
||||||
|
f'|__msg_spec__:\n'
|
||||||
|
f'{speclines}\n',
|
||||||
|
prefix=' '*2,
|
||||||
|
)
|
||||||
|
return (
|
||||||
|
f'<{type(self).__name__}(\n'
|
||||||
|
f'{body}'
|
||||||
|
')>'
|
||||||
|
)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def pld_spec_str(self) -> str:
|
def pld_spec_str(self) -> str:
|
||||||
spec: Union[Type]|Type = self.pld_spec
|
spec: Union[Type]|Type = self.pld_spec
|
||||||
|
@ -163,8 +187,8 @@ class MsgCodec(Struct):
|
||||||
|
|
||||||
) -> bytes:
|
) -> bytes:
|
||||||
'''
|
'''
|
||||||
Encode input python objects to `msgpack` bytes for transfer
|
Encode input python objects to `msgpack` bytes for
|
||||||
on a tranport protocol connection.
|
transfer on a tranport protocol connection.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
return self._enc.encode(py_obj)
|
return self._enc.encode(py_obj)
|
||||||
|
@ -325,15 +349,9 @@ class MsgCodec(Struct):
|
||||||
|
|
||||||
|
|
||||||
def mk_codec(
|
def mk_codec(
|
||||||
ipc_msg_spec: Union[Type[Struct]]|Any|None = None,
|
|
||||||
#
|
|
||||||
# ^TODO^: in the long run, do we want to allow using a diff IPC `Msg`-set?
|
|
||||||
# it would break the runtime, but maybe say if you wanted
|
|
||||||
# to add some kinda field-specific or wholesale `.pld` ecryption?
|
|
||||||
|
|
||||||
# struct type unions set for `Decoder`
|
# struct type unions set for `Decoder`
|
||||||
# https://jcristharif.com/msgspec/structs.html#tagged-unions
|
# https://jcristharif.com/msgspec/structs.html#tagged-unions
|
||||||
ipc_pld_spec: Union[Type[Struct]]|Any|None = None,
|
ipc_pld_spec: Union[Type[Struct]]|Any = Any,
|
||||||
|
|
||||||
# TODO: offering a per-msg(-field) type-spec such that
|
# TODO: offering a per-msg(-field) type-spec such that
|
||||||
# the fields can be dynamically NOT decoded and left as `Raw`
|
# the fields can be dynamically NOT decoded and left as `Raw`
|
||||||
|
@ -352,7 +370,6 @@ def mk_codec(
|
||||||
dec_hook: Callable|None = None,
|
dec_hook: Callable|None = None,
|
||||||
enc_hook: Callable|None = None,
|
enc_hook: Callable|None = None,
|
||||||
# ------ - ------
|
# ------ - ------
|
||||||
**kwargs,
|
|
||||||
#
|
#
|
||||||
# Encoder:
|
# Encoder:
|
||||||
# write_buffer_size=write_buffer_size,
|
# write_buffer_size=write_buffer_size,
|
||||||
|
@ -367,28 +384,6 @@ def mk_codec(
|
||||||
`msgspec` ;).
|
`msgspec` ;).
|
||||||
|
|
||||||
'''
|
'''
|
||||||
if (
|
|
||||||
ipc_msg_spec is not None
|
|
||||||
and ipc_pld_spec
|
|
||||||
):
|
|
||||||
raise RuntimeError(
|
|
||||||
f'If a payload spec is provided,\n'
|
|
||||||
"the builtin SC-shuttle-protocol's msg set\n"
|
|
||||||
f'(i.e. a `{MsgType}`) MUST be used!\n\n'
|
|
||||||
f'However both values were passed as => mk_codec(\n'
|
|
||||||
f' ipc_msg_spec={ipc_msg_spec}`\n'
|
|
||||||
f' ipc_pld_spec={ipc_pld_spec}`\n)\n'
|
|
||||||
)
|
|
||||||
|
|
||||||
elif (
|
|
||||||
ipc_pld_spec
|
|
||||||
and
|
|
||||||
|
|
||||||
# XXX required for now (or maybe forever?) until
|
|
||||||
# we can dream up a way to allow parameterizing and/or
|
|
||||||
# custom overrides to the `Msg`-spec protocol itself?
|
|
||||||
ipc_msg_spec is None
|
|
||||||
):
|
|
||||||
# (manually) generate a msg-payload-spec for all relevant
|
# (manually) generate a msg-payload-spec for all relevant
|
||||||
# god-boxing-msg subtypes, parameterizing the `Msg.pld: PayloadT`
|
# god-boxing-msg subtypes, parameterizing the `Msg.pld: PayloadT`
|
||||||
# for the decoder such that all sub-type msgs in our SCIPP
|
# for the decoder such that all sub-type msgs in our SCIPP
|
||||||
|
@ -403,9 +398,6 @@ def mk_codec(
|
||||||
assert len(ipc_msg_spec.__args__) == len(msg_types)
|
assert len(ipc_msg_spec.__args__) == len(msg_types)
|
||||||
assert ipc_msg_spec
|
assert ipc_msg_spec
|
||||||
|
|
||||||
else:
|
|
||||||
ipc_msg_spec = ipc_msg_spec or Any
|
|
||||||
|
|
||||||
enc = msgpack.Encoder(
|
enc = msgpack.Encoder(
|
||||||
enc_hook=enc_hook,
|
enc_hook=enc_hook,
|
||||||
)
|
)
|
||||||
|
@ -418,8 +410,6 @@ def mk_codec(
|
||||||
_enc=enc,
|
_enc=enc,
|
||||||
_dec=dec,
|
_dec=dec,
|
||||||
pld_spec=ipc_pld_spec,
|
pld_spec=ipc_pld_spec,
|
||||||
# payload_msg_specs=payload_msg_specs,
|
|
||||||
# **kwargs,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# sanity on expected backend support
|
# sanity on expected backend support
|
||||||
|
@ -500,8 +490,16 @@ def apply_codec(
|
||||||
- https://github.com/oremanj/tricycle/blob/master/tricycle/_tests/test_tree_var.py
|
- https://github.com/oremanj/tricycle/blob/master/tricycle/_tests/test_tree_var.py
|
||||||
|
|
||||||
'''
|
'''
|
||||||
|
__tracebackhide__: bool = True
|
||||||
orig: MsgCodec = _ctxvar_MsgCodec.get()
|
orig: MsgCodec = _ctxvar_MsgCodec.get()
|
||||||
assert orig is not codec
|
assert orig is not codec
|
||||||
|
if codec.pld_spec is None:
|
||||||
|
breakpoint()
|
||||||
|
|
||||||
|
log.info(
|
||||||
|
'Applying new msg-spec codec\n\n'
|
||||||
|
f'{codec}\n'
|
||||||
|
)
|
||||||
token: RunVarToken = _ctxvar_MsgCodec.set(codec)
|
token: RunVarToken = _ctxvar_MsgCodec.set(codec)
|
||||||
|
|
||||||
# TODO: for TreeVar approach, see docs for @cm `.being()` API:
|
# TODO: for TreeVar approach, see docs for @cm `.being()` API:
|
||||||
|
@ -518,7 +516,10 @@ def apply_codec(
|
||||||
_ctxvar_MsgCodec.reset(token)
|
_ctxvar_MsgCodec.reset(token)
|
||||||
|
|
||||||
assert _ctxvar_MsgCodec.get() is orig
|
assert _ctxvar_MsgCodec.get() is orig
|
||||||
|
log.info(
|
||||||
|
'Reverted to last msg-spec codec\n\n'
|
||||||
|
f'{orig}\n'
|
||||||
|
)
|
||||||
|
|
||||||
def current_codec() -> MsgCodec:
|
def current_codec() -> MsgCodec:
|
||||||
'''
|
'''
|
||||||
|
@ -532,14 +533,15 @@ def current_codec() -> MsgCodec:
|
||||||
|
|
||||||
@cm
|
@cm
|
||||||
def limit_msg_spec(
|
def limit_msg_spec(
|
||||||
payload_types: Union[Type[Struct]],
|
payload_spec: Union[Type[Struct]],
|
||||||
|
|
||||||
# TODO: don't need this approach right?
|
# TODO: don't need this approach right?
|
||||||
# -> related to the `MsgCodec._payload_decs` stuff above..
|
# -> related to the `MsgCodec._payload_decs` stuff above..
|
||||||
# tagged_structs: list[Struct]|None = None,
|
# tagged_structs: list[Struct]|None = None,
|
||||||
|
|
||||||
**codec_kwargs,
|
**codec_kwargs,
|
||||||
):
|
|
||||||
|
) -> MsgCodec:
|
||||||
'''
|
'''
|
||||||
Apply a `MsgCodec` that will natively decode the SC-msg set's
|
Apply a `MsgCodec` that will natively decode the SC-msg set's
|
||||||
`Msg.pld: Union[Type[Struct]]` payload fields using
|
`Msg.pld: Union[Type[Struct]]` payload fields using
|
||||||
|
@ -547,10 +549,37 @@ def limit_msg_spec(
|
||||||
for all IPC contexts in use by the current `trio.Task`.
|
for all IPC contexts in use by the current `trio.Task`.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
|
__tracebackhide__: bool = True
|
||||||
|
curr_codec = current_codec()
|
||||||
msgspec_codec: MsgCodec = mk_codec(
|
msgspec_codec: MsgCodec = mk_codec(
|
||||||
payload_types=payload_types,
|
ipc_pld_spec=payload_spec,
|
||||||
**codec_kwargs,
|
**codec_kwargs,
|
||||||
)
|
)
|
||||||
with apply_codec(msgspec_codec) as applied_codec:
|
with apply_codec(msgspec_codec) as applied_codec:
|
||||||
assert applied_codec is msgspec_codec
|
assert applied_codec is msgspec_codec
|
||||||
yield msgspec_codec
|
yield msgspec_codec
|
||||||
|
|
||||||
|
assert curr_codec is current_codec()
|
||||||
|
|
||||||
|
|
||||||
|
# XXX: msgspec won't allow this with non-struct custom types
|
||||||
|
# like `NamespacePath`!@!
|
||||||
|
# @cm
|
||||||
|
# def extend_msg_spec(
|
||||||
|
# payload_spec: Union[Type[Struct]],
|
||||||
|
|
||||||
|
# ) -> MsgCodec:
|
||||||
|
# '''
|
||||||
|
# Extend the current `MsgCodec.pld_spec` (type set) by extending
|
||||||
|
# the payload spec to **include** the types specified by
|
||||||
|
# `payload_spec`.
|
||||||
|
|
||||||
|
# '''
|
||||||
|
# codec: MsgCodec = current_codec()
|
||||||
|
# pld_spec: Union[Type] = codec.pld_spec
|
||||||
|
# extended_spec: Union[Type] = pld_spec|payload_spec
|
||||||
|
|
||||||
|
# with limit_msg_spec(payload_types=extended_spec) as ext_codec:
|
||||||
|
# # import pdbp; pdbp.set_trace()
|
||||||
|
# assert ext_codec.pld_spec == extended_spec
|
||||||
|
# yield ext_codec
|
||||||
|
|
Loading…
Reference in New Issue