Do proper `wrapt` arg extraction for type checking

Use an inner function / closure to properly process required arguments
at call time as is recommended in the `wrap` docs. Do async gen and
arg introspection at decorate time and raise appropriate type errors.
contexts
Tyler Goodlet 2019-01-25 00:10:13 -05:00
parent 1b405ab4fe
commit 3d0de25f93
1 changed files with 71 additions and 48 deletions

View File

@ -1,8 +1,9 @@
""" """
Messaging pattern APIs and helpers. Messaging pattern APIs and helpers.
""" """
import inspect
import typing import typing
from typing import Dict, Any, Set, Union from typing import Dict, Any, Set, Union, Callable
from functools import partial from functools import partial
from async_generator import aclosing from async_generator import aclosing
@ -11,6 +12,7 @@ import wrapt
from .log import get_logger from .log import get_logger
from . import current_actor from . import current_actor
from ._ipc import Context
__all__ = ['pub'] __all__ = ['pub']
@ -181,19 +183,21 @@ def pub(
@wrapt.decorator(adapter=takes_ctx) @wrapt.decorator(adapter=takes_ctx)
async def wrapper(agen, instance, args, kwargs): async def wrapper(agen, instance, args, kwargs):
task_name = None # this is used to extract arguments properly as per
if tasks: # the `wrapt` docs
try: async def _execute(
task_name = kwargs.pop('task_name') ctx: Context,
except KeyError: topics: Set[str],
*args,
# *,
task_name: str = None,
packetizer: Callable = None,
**kwargs,
):
if tasks and task_name is None:
raise TypeError( raise TypeError(
f"{agen} must be called with a `task_name` named argument " f"{agen} must be called with a `task_name` named "
f"with a falue from {tasks}") f"argument with a falue from {tasks}")
# pop required kwargs used internally
ctx = kwargs.pop('ctx')
topics = kwargs.pop('topics')
packetizer = kwargs.pop('packetizer', None)
ss = current_actor().statespace ss = current_actor().statespace
lockmap = ss.setdefault('_pubtask2lock', task2lock) lockmap = ss.setdefault('_pubtask2lock', task2lock)
@ -212,16 +216,19 @@ def pub(
respawn = True respawn = True
while respawn: while respawn:
respawn = False respawn = False
log.info(f"Spawning data feed task for {agen.__name__}") log.info(
f"Spawning data feed task for {funcname}")
try: try:
# unblocks when no more symbols subscriptions exist # unblocks when no more symbols subscriptions exist
# and the streamer task terminates # and the streamer task terminates
await fan_out_to_ctxs( await fan_out_to_ctxs(
pub_async_gen_func=partial(agen, *args, **kwargs), pub_async_gen_func=partial(
agen, *args, **kwargs),
topics2ctxs=topics2ctxs, topics2ctxs=topics2ctxs,
packetizer=packetizer, packetizer=packetizer,
) )
log.info(f"Terminating stream task {task_name or ''}" log.info(
f"Terminating stream task {task_name or ''}"
f" for {agen.__name__}") f" for {agen.__name__}")
except trio.BrokenResourceError: except trio.BrokenResourceError:
log.exception("Respawning failed data feed task") log.exception("Respawning failed data feed task")
@ -233,6 +240,22 @@ def pub(
# if there are truly no more subscriptions with this broker # if there are truly no more subscriptions with this broker
# drop from broker subs dict # drop from broker subs dict
if not any(topics2ctxs.values()): if not any(topics2ctxs.values()):
log.info(f"No more subscriptions for publisher {task_name}") log.info(
f"No more subscriptions for publisher {task_name}")
# invoke it
await _execute(*args, **kwargs)
funcname = wrapped.__name__
if not inspect.isasyncgenfunction(wrapped):
raise TypeError(
f"Publisher {funcname} must be an async generator function"
)
if 'get_topics' not in inspect.signature(wrapped).parameters:
raise TypeError(
f"Publisher async gen {funcname} must define a "
"`get_topics` argument"
)
return wrapper(wrapped) return wrapper(wrapped)