jc211
/
tractor
forked from goodboy/tractor
1
0
Fork 0
Code Pull Requests Activity

Proxy through `dec_hook` in `.limit_plds()` APIs

runtime_to_msgspec
Tyler Goodlet 2024-06-17 09:23:31 -04:00
parent a0ee0cc713
commit 04bd111037
1 changed files with 20 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
tractor/msg/_ops.py
View File

@ -27,6 +27,7 @@ from contextlib import (
) )
from typing import ( from typing import (
Any, Any,
Callable,
Type, Type,
TYPE_CHECKING, TYPE_CHECKING,
Union, Union,
@ -138,6 +139,7 @@ class PldRx(Struct):
def limit_plds( def limit_plds(
self, self,
spec: Union[Type[Struct]], spec: Union[Type[Struct]],
**dec_kwargs,
) -> MsgDec: ) -> MsgDec:
''' '''
@ -147,7 +149,10 @@ class PldRx(Struct):
''' '''
orig_dec: MsgDec = self._pld_dec orig_dec: MsgDec = self._pld_dec
limit_dec: MsgDec = mk_dec(spec=spec) limit_dec: MsgDec = mk_dec(
spec=spec,
**dec_kwargs,
)
try: try:
self._pld_dec = limit_dec self._pld_dec = limit_dec
yield limit_dec yield limit_dec
@ -449,7 +454,7 @@ class PldRx(Struct):
@cm @cm
def limit_plds( def limit_plds(
spec: Union[Type[Struct]], spec: Union[Type[Struct]],
**kwargs, **dec_kwargs,
) -> MsgDec: ) -> MsgDec:
''' '''
@ -467,7 +472,7 @@ def limit_plds(
with rx.limit_plds( with rx.limit_plds(
spec=spec, spec=spec,
**kwargs, **dec_kwargs,
) as pldec: ) as pldec:
log.runtime( log.runtime(
'Applying payload-decoder\n\n' 'Applying payload-decoder\n\n'
@ -487,7 +492,9 @@ def limit_plds(
async def maybe_limit_plds( async def maybe_limit_plds(
ctx: Context, ctx: Context,
spec: Union[Type[Struct]]|None = None, spec: Union[Type[Struct]]|None = None,
dec_hook: Callable|None = None,
**kwargs, **kwargs,
) -> MsgDec|None: ) -> MsgDec|None:
''' '''
Async compat maybe-payload type limiter. Async compat maybe-payload type limiter.
@ -497,7 +504,11 @@ async def maybe_limit_plds(
used. used.
''' '''
if spec is None: if (
spec is None
and
dec_hook is None
):
yield None yield None
return return
@ -505,7 +516,11 @@ async def maybe_limit_plds(
curr_ctx: Context = current_ipc_ctx() curr_ctx: Context = current_ipc_ctx()
assert ctx is curr_ctx assert ctx is curr_ctx
with ctx._pld_rx.limit_plds(spec=spec) as msgdec: with ctx._pld_rx.limit_plds(
spec=spec,
dec_hook=dec_hook,
**kwargs,
) as msgdec:
yield msgdec yield msgdec
curr_ctx: Context = current_ipc_ctx() curr_ctx: Context = current_ipc_ctx()